From 8b4d93ba2b4e4e79062fd59a35a623e381dcf6b0 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:09:00 -0400 Subject: [PATCH 001/138] [Perf] Remove redundant clone for GLM, Deepseek etc (#46651) Signed-off-by: yewentao256 --- vllm/model_executor/models/AXK1.py | 2 +- vllm/model_executor/models/deepseek_v2.py | 2 +- vllm/model_executor/models/glm4_moe_lite.py | 2 +- vllm/model_executor/models/openpangu.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 701ec67c855..d526f57d3d9 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -649,7 +649,7 @@ class AXK1DecoderLayer(nn.Module): ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention if residual is None: - residual = hidden_states.clone() + residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 2f6a472fe35..8d20e0b5c68 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1186,7 +1186,7 @@ class DeepseekV2DecoderLayer(nn.Module): ) -> torch.Tensor: # Self Attention if residual is None: - residual = hidden_states.clone() + residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index 77aaa179aa5..b4d0fe96680 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -184,7 +184,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module): ) -> torch.Tensor: # Self Attention if residual is None: - residual = hidden_states.clone() + residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index a517c52e690..8432566a150 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -935,7 +935,7 @@ class OpenPanguDecoderLayer(nn.Module): residual: torch.Tensor | None, ) -> torch.Tensor: if residual is None: - residual = hidden_states.clone() + residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) From c5e3c40877c2b6d0e16d534641b39fe6744979b7 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:13:08 -0400 Subject: [PATCH 002/138] Fix P/D with DP Supervisor (#46628) Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- vllm/v1/engine/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 57a788631ce..f97f697dedc 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1175,10 +1175,10 @@ class EngineCoreProc(EngineCore): numa_utils.log_current_affinity_state(process_title) if data_parallel and vllm_config.kv_transfer_config is not None: - # modify the engine_id and append the local_dp_rank to it to ensure + # modify the engine_id and append the dp_rank to it to ensure # that the kv_transfer_config is unique for each DP rank. vllm_config.kv_transfer_config.engine_id = ( - f"{vllm_config.kv_transfer_config.engine_id}_dp{local_dp_rank}" + f"{vllm_config.kv_transfer_config.engine_id}_dp{dp_rank}" ) logger.debug( "Setting kv_transfer_config.engine_id to %s", From 2a6f8f0c05ab1dd0b11540157fb72b6888883aab Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 25 Jun 2026 15:24:09 -0500 Subject: [PATCH 003/138] [ROCm][CI] Fine-tuning queues and test names (#39238) Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 30 ++++++--------------- .buildkite/test_areas/cuda.yaml | 4 +-- .buildkite/test_areas/distributed.yaml | 14 +++++----- .buildkite/test_areas/e2e_integration.yaml | 12 ++++----- .buildkite/test_areas/kernels.yaml | 8 +++--- .buildkite/test_areas/lm_eval.yaml | 31 +++++++++++----------- 6 files changed, 43 insertions(+), 56 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index eeb685e9892..09d8a33cc3f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -771,7 +771,7 @@ steps: #----------------------------------------------------------- mi300 ยท cuda ------------------------------------------------------------# -- label: Platform Tests (CUDA) # TBD +- label: Platform Tests # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 @@ -831,7 +831,7 @@ steps: - pytest -v -s distributed/test_eplb_execute.py - pytest -v -s distributed/test_eplb_spec_decode.py -- label: Distributed Tests (2xH100-2xMI250) # TBD +- label: Distributed Tests (2xH100-2xMI300) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 @@ -843,13 +843,19 @@ steps: - vllm/model_executor/layers/fused_moe/ - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py + - tests/v1/distributed/test_dbo.py - tests/distributed/test_context_parallel.py - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - pytest -v -s tests/distributed/test_context_parallel.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.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 + - 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 (4xA100-4xMI300) # TBD timeout_in_minutes: 180 @@ -2195,26 +2201,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 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: - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.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: Metrics, Tracing (2 GPUs) # TBD timeout_in_minutes: 20 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index b56e635bea6..956c76cf05f 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -2,8 +2,8 @@ group: CUDA depends_on: - image-build steps: -- label: Platform Tests (CUDA) - key: platform-tests-cuda +- label: Platform Tests + key: platform-tests timeout_in_minutes: 15 device: h200_18gb source_file_dependencies: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index b880fadf356..5ff4b24b744 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -174,8 +174,8 @@ steps: # test multi-node TP with multiproc executor (simulated on single node) - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node -- label: Distributed Tests (8 GPUs)(H100) - key: distributed-tests-8-gpus-h100 +- label: Distributed Tests (8xH100) + key: distributed-tests-8xh100 timeout_in_minutes: 10 device: h100 num_devices: 8 @@ -195,8 +195,8 @@ steps: # test with torchrun tp=2 and dp=4 with ep - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep -- label: Distributed Tests (4 GPUs)(A100) - key: distributed-tests-4-gpus-a100 +- label: Distributed Tests (4xA100) + key: distributed-tests-4xa100 device: a100 optional: true num_devices: 4 @@ -211,7 +211,7 @@ steps: - pytest -v -s -x lora/test_mixtral.py - label: Distributed Tests (2xH100-2xMI300) - key: distributed-tests-2-gpus-h100 + key: distributed-tests-2xh100-2xmi300 timeout_in_minutes: 15 device: h100 optional: true @@ -237,8 +237,8 @@ steps: - 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) - key: distributed-tests-2-gpus-b200 +- label: Distributed Tests (2xB200) + key: distributed-tests-2xb200 device: b200-k8s optional: true working_dir: "/vllm-workspace/" diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 88039a33960..3f87e3958d0 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -2,8 +2,8 @@ group: E2E Integration depends_on: - image-build steps: -- label: DeepSeek V2-Lite Sync EPLB Accuracy - key: deepseek-v2-lite-sync-eplb-accuracy +- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100) + key: deepseek-v2-lite-sync-eplb-accuracy-4xh100 timeout_in_minutes: 60 device: h100 optional: true @@ -12,8 +12,8 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy - key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100) + key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100 timeout_in_minutes: 60 device: h100 optional: true @@ -22,8 +22,8 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200) - key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200 +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (2xB200) + key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200 timeout_in_minutes: 60 device: b200-k8s optional: true diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index c5341a0f518..10c132da095 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -274,8 +274,8 @@ steps: - pytest -v -s kernels/helion/ -- label: Kernels FP8 MoE Test (1 H100) - key: kernels-fp8-moe-test-1-h100 +- label: Kernels FP8 MoE Test (1xH100) + key: kernels-fp8-moe-test-1xh100 timeout_in_minutes: 90 device: h100 num_devices: 1 @@ -291,8 +291,8 @@ steps: - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py -- label: Kernels FP8 MoE Test (2 H100s) - key: kernels-fp8-moe-test-2-h100s +- label: Kernels FP8 MoE Test (2xH100) + key: kernels-fp8-moe-test-2xh100 timeout_in_minutes: 90 device: h100 num_devices: 2 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index a64edbd1c4f..d5c4b6957ab 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -28,7 +28,8 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py -# - label: LM Eval Large Models (4 GPUs)(A100) +# - label: LM Eval Large Models (4xA100) +# key: lm-eval-large-models-4xa100 # device: a100 # optional: true # num_devices: 4 @@ -40,8 +41,8 @@ steps: # - 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: LM Eval Large Models (4 GPUs)(H100) - key: lm-eval-large-models-4-gpus-h100 +- label: LM Eval Large Models (4xH100) + key: lm-eval-large-models-4xh100 device: h100 optional: true num_devices: 4 @@ -53,8 +54,8 @@ steps: - export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4 -- label: LM Eval Small Models (B200) - key: lm-eval-small-models-b200 +- label: LM Eval Small Models (2xB200) + key: lm-eval-small-models-2xb200 timeout_in_minutes: 120 device: b200-k8s optional: true @@ -64,8 +65,8 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt -- label: LM Eval Large Models (B200, EP) - key: lm-eval-large-models-b200-ep +- label: LM Eval Large Models EP (2xB200) + key: lm-eval-large-models-ep-2xb200 timeout_in_minutes: 120 device: b200-k8s optional: true @@ -76,8 +77,8 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt -- label: LM Eval Qwen3.5 Models (B200) - key: lm-eval-qwen3-5-models-b200 +- label: LM Eval Qwen3.5 Models (2xB200) + key: lm-eval-qwen3-5-models-2xb200 timeout_in_minutes: 120 device: b200-k8s optional: true @@ -93,8 +94,8 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt -- label: LM Eval Large Models (H200) - key: lm-eval-large-models-h200 +- label: LM Eval Large Models (8xH200) + key: lm-eval-large-models-8xh200 timeout_in_minutes: 60 device: h200 optional: true @@ -192,8 +193,8 @@ steps: 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) - key: gpqa-eval-gpt-oss-h100 +- label: GPQA Eval (GPT-OSS) (2xH100) + key: gpqa-eval-gpt-oss-2xh100 timeout_in_minutes: 120 device: h100 optional: true @@ -206,8 +207,8 @@ steps: - 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-h100.txt -- label: GPQA Eval (GPT-OSS) (B200) - key: gpqa-eval-gpt-oss-b200 +- label: GPQA Eval (GPT-OSS) (2xB200) + key: gpqa-eval-gpt-oss-2xb200 timeout_in_minutes: 120 device: b200-k8s optional: true From e8c24a769576fa318ca93fbd927128246a534325 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 25 Jun 2026 15:02:28 -0600 Subject: [PATCH 004/138] [Kernel] Vectorized fp32 `moe_sum` reduction and support any topk (#46643) Signed-off-by: mgoin Co-authored-by: Claude --- .../moe/moe_align_sum_kernels.cu | 204 ++++++++++++++---- tests/kernels/moe/test_moe.py | 20 +- 2 files changed, 175 insertions(+), 49 deletions(-) diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index 1e842381349..985c47b0765 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -11,6 +11,7 @@ #include "../../cuda_compat.h" #include "libtorch_stable/core/math.hpp" #include "libtorch_stable/dispatch_utils.h" +#include "libtorch_stable/quantization/vectorization.cuh" #include "libtorch_stable/torch_utils.h" #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) @@ -349,19 +350,102 @@ __global__ void count_and_sort_expert_tokens_kernel( max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map); } +// Reduce the topk expert outputs per token (summed in fp32). The output is +// dense [num_tokens, d]; the input is addressed by its strides so non- +// contiguous inputs work without a copy. A 16B-vectorized path is used when +// the hidden dim is contiguous (innermost stride 1) and aligned; otherwise a +// scalar kernel reads via arbitrary strides. topk is a compile-time constant +// for common values and runtime otherwise. + +// Elements per 16-byte vector (8 for bf16/fp16, 4 for fp32). +template +constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t); + template -__global__ void moe_sum_kernel( - scalar_t* __restrict__ out, // [..., d] - const scalar_t* __restrict__ input, // [..., topk, d] - const int d) { - const int64_t token_idx = blockIdx.x; - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { - scalar_t x = 0.0; +__global__ void moe_sum_vec_kernel( + scalar_t* __restrict__ out, // [num_tokens, d], contiguous + const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous + const int64_t num_tokens, const int d, const int64_t stride_token, + const int64_t stride_topk) { + using vec_t = vllm::vec_n_t>; // 16-byte pack + constexpr int VEC = MOE_SUM_VEC; + const int64_t n_vec = d / VEC; + const int64_t total = num_tokens * n_vec; + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; + i += (int64_t)gridDim.x * blockDim.x) { + const int64_t token = i / n_vec; + const int64_t v = i % n_vec; + const scalar_t* in_tok = input + token * stride_token + v * VEC; + + float acc[VEC]; +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] = 0.f; + #pragma unroll for (int k = 0; k < TOPK; ++k) { - x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]); + vec_t packed = *reinterpret_cast(in_tok + k * stride_topk); +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] += static_cast(packed.val[j]); } - out[token_idx * d + idx] = x; + + vec_t outp; +#pragma unroll + for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast(acc[j]); + *reinterpret_cast(out + token * d + v * VEC) = outp; + } +} + +// Runtime-topk variant of the above. +template +__global__ void moe_sum_vec_dynamic_kernel( + scalar_t* __restrict__ out, // [num_tokens, d], contiguous + const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous + const int64_t num_tokens, const int d, const int topk, + const int64_t stride_token, const int64_t stride_topk) { + using vec_t = vllm::vec_n_t>; + constexpr int VEC = MOE_SUM_VEC; + const int64_t n_vec = d / VEC; + const int64_t total = num_tokens * n_vec; + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; + i += (int64_t)gridDim.x * blockDim.x) { + const int64_t token = i / n_vec; + const int64_t v = i % n_vec; + const scalar_t* in_tok = input + token * stride_token + v * VEC; + + float acc[VEC]; +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] = 0.f; + + for (int k = 0; k < topk; ++k) { + vec_t packed = *reinterpret_cast(in_tok + k * stride_topk); +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] += static_cast(packed.val[j]); + } + + vec_t outp; +#pragma unroll + for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast(acc[j]); + *reinterpret_cast(out + token * d + v * VEC) = outp; + } +} + +// Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims +// (including a non-contiguous hidden stride) via per-element strided reads. +template +__global__ void moe_sum_scalar_kernel( + scalar_t* __restrict__ out, // [num_tokens, d], contiguous + const scalar_t* __restrict__ input, // [num_tokens, topk, d] + const int d, const int topk, const int64_t stride_token, + const int64_t stride_topk, const int64_t stride_hidden) { + const int64_t token_idx = blockIdx.x; + const scalar_t* in_tok = input + token_idx * stride_token; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + float x = 0.f; + for (int k = 0; k < topk; ++k) { + x += static_cast( + VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden])); + } + out[token_idx * d + idx] = static_cast(x); } } @@ -626,52 +710,82 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch, void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size] torch::stable::Tensor& output) // [num_tokens, hidden_size] { + // Output is dense and written in place, so it must be contiguous. The input + // is read by its strides (no copy); only the hidden dim needs to be + // contiguous to take the vectorized path. + STD_TORCH_CHECK(output.is_contiguous(), + "moe_sum expects a contiguous output"); + const int hidden_size = input.size(-1); - const auto num_tokens = output.numel() / hidden_size; + const int64_t num_tokens = output.numel() / hidden_size; const int topk = input.size(1); + const int64_t stride_token = input.stride(0); + const int64_t stride_topk = input.stride(1); + const int64_t stride_hidden = input.stride(2); - dim3 grid(num_tokens); - dim3 block(std::min(hidden_size, 1024)); const torch::stable::accelerator::DeviceGuard device_guard( output.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(output.get_device_index()); - switch (topk) { - case 2: - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - input.scalar_type(), "moe_sum_kernel", [&] { - vllm::moe::moe_sum_kernel<<>>( - reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast(input.const_data_ptr()), - hidden_size); - }); - break; +#define LAUNCH_MOE_SUM_VEC(TOPK) \ + vllm::moe::moe_sum_vec_kernel \ + <<>>( \ + out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk) - case 3: - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - input.scalar_type(), "moe_sum_kernel", [&] { - vllm::moe::moe_sum_kernel<<>>( - reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast(input.const_data_ptr()), - hidden_size); - }); - break; + VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] { + constexpr int VEC = vllm::moe::MOE_SUM_VEC; + constexpr int WIDTH = VEC * sizeof(scalar_t); // 16 bytes + auto* out_ptr = reinterpret_cast(output.mutable_data_ptr()); + auto* in_ptr = reinterpret_cast(input.const_data_ptr()); - case 4: - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - input.scalar_type(), "moe_sum_kernel", [&] { - vllm::moe::moe_sum_kernel<<>>( - reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast(input.const_data_ptr()), - hidden_size); - }); - break; - - default: - torch::stable::sum_out(output, input, std::array{1}); - break; - } + // Vectorize along hidden only when it is contiguous (innermost stride 1), + // a whole number of vectors, and every row offset stays 16B-aligned. + const bool can_vec = (stride_hidden == 1) && (hidden_size % VEC == 0) && + (stride_token % VEC == 0) && + (stride_topk % VEC == 0) && + (reinterpret_cast(in_ptr) % WIDTH == 0) && + (reinterpret_cast(out_ptr) % WIDTH == 0); + if (can_vec) { + const int64_t n_vec = hidden_size / VEC; + const int64_t total = num_tokens * n_vec; + const int block = 256; + const dim3 grid(std::min((total + block - 1) / block, 65535)); + switch (topk) { + case 1: + LAUNCH_MOE_SUM_VEC(1); + break; + case 2: + LAUNCH_MOE_SUM_VEC(2); + break; + case 4: + LAUNCH_MOE_SUM_VEC(4); + break; + case 6: + LAUNCH_MOE_SUM_VEC(6); + break; + case 8: + LAUNCH_MOE_SUM_VEC(8); + break; + case 9: + LAUNCH_MOE_SUM_VEC(9); + break; + default: + vllm::moe::moe_sum_vec_dynamic_kernel + <<>>(out_ptr, in_ptr, num_tokens, + hidden_size, topk, + stride_token, stride_topk); + break; + } + } else { + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + vllm::moe::moe_sum_scalar_kernel<<>>( + out_ptr, in_ptr, hidden_size, topk, stride_token, stride_topk, + stride_hidden); + } + }); +#undef LAUNCH_MOE_SUM_VEC } void moe_lora_align_block_size( diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 45cd17b3b11..f8b98c82a24 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1243,15 +1243,27 @@ def test_batched_moe_align_block_size_opcheck(): ) +# topk=8 covers topk > 4; k=511 covers the non-vectorized scalar path. The +# layouts exercise contiguous input plus the two non-contiguous cases: a +# transpose (strided hidden -> scalar gather) and a topk-slice (hidden still +# contiguous -> vectorized). @pytest.mark.parametrize("m", [1, 33, 222]) -@pytest.mark.parametrize("topk", TOP_KS) +@pytest.mark.parametrize("topk", [*TOP_KS, 8]) @pytest.mark.parametrize("k", [128, 511, 1024]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype): - input = torch.randn((m, topk, k), device="cuda", dtype=dtype) +@pytest.mark.parametrize("layout", ["contig", "transpose", "slice"]) +def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype, layout: str): + if layout == "transpose": + input = torch.randn((m, k, topk), device="cuda", dtype=dtype).transpose(1, 2) + elif layout == "slice": + input = torch.randn((m, 2 * topk, k), device="cuda", dtype=dtype)[:, ::2, :] + else: + input = torch.randn((m, topk, k), device="cuda", dtype=dtype) + assert input.is_contiguous() == (layout == "contig") actual = torch.empty((m, k), device="cuda", dtype=dtype) - expected = input.sum(dim=1) + # Reduction accumulates in fp32. + expected = input.float().sum(dim=1).to(dtype) torch.ops._moe_C.moe_sum(input, actual) torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0) From a2e8ec3d52ab4e163501c8c7bee8c03ca8359a7a Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 25 Jun 2026 15:07:04 -0600 Subject: [PATCH 005/138] [CI] Depend GPQA Eval DGX Spark job on arm64 image build (#46736) Signed-off-by: mgoin --- .buildkite/test_areas/lm_eval.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d5c4b6957ab..8063d5e72fd 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -227,6 +227,8 @@ steps: device: dgx-spark optional: true num_devices: 1 + depends_on: + - arm64-image-build source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization From 27da2a2ac4776faa4265cba38ba86dd3a7119c4f Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Thu, 25 Jun 2026 17:08:04 -0500 Subject: [PATCH 006/138] [Hardware][AMD][CI] Use Triton-based AITER MHA for LM Eval Qwen-3.5 Models Tests (#46691) Signed-off-by: Matthew Wong --- .buildkite/test-amd.yaml | 7 +++---- .../gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml | 2 +- vllm/v1/attention/ops/vit_attn_wrappers.py | 6 +++++- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 09d8a33cc3f..3a6568c2e0a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2790,13 +2790,12 @@ steps: - 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) %N # TBD - timeout_in_minutes: 180 +- 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 - parallelism: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/model_executor/models/qwen3_5.py @@ -2811,7 +2810,7 @@ steps: - 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 --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - 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 diff --git a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml index 2c0431747d0..ca5cc450c07 100644 --- a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml +++ b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml @@ -3,7 +3,6 @@ accuracy_threshold: 0.89 tolerance: 0.03 num_questions: 1319 num_fewshot: 5 -startup_max_wait_seconds: 3600 server_args: >- --max-model-len 4096 --tensor-parallel-size 2 @@ -11,3 +10,4 @@ server_args: >- --moe-backend aiter env: VLLM_ROCM_USE_AITER: "1" + ENABLE_CK: "0" # Avoid AITER CK-based MHA JIT compilation to save time diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index 4506f452cf9..5bbcc3386e5 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -12,6 +12,8 @@ latencies by ~7% (see qwen2_5_vl for example usage) To use these ops, you must have a recent version of PyTorch installed (>= 2.4.0) """ +from typing import Any + import einops import torch import torch.nn.functional as F @@ -31,9 +33,11 @@ def flash_attn_maxseqlen_wrapper( cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: - kwargs = {} + kwargs: dict[str, Any] = {} if is_rocm_aiter: from aiter import flash_attn_varlen_func + + kwargs["window_size"] = (-1, -1) else: from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func From c53994e1348bac3496aafb88e9e731124a00a8a7 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:46:10 -0500 Subject: [PATCH 007/138] [Model Runner V2][Spec Decode] Use log1p to compute residual during rejection sampling (#46665) Signed-off-by: Giancarlo Delfin --- .../gpu/spec_decode/rejection_sampler_utils.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index bad70aa0451..7020f228046 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import tl, triton +from vllm.triton_utils import tl, tldevice, triton from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64 @@ -387,14 +387,16 @@ def _resample_kernel( draft_lse = tl.load(draft_rejected_logsumexp_ptr + req_idx) target_log_probs = target_logits - target_lse draft_log_probs = draft_logits - draft_lse - # Compute the residual: max(p(x) - q(x), 0) - # Equivalent log form: log(max(exp(log_p(x)) - exp(log_q(x)), 0)) + # Compute the residual: + # r(x) = max(p(x) - q(x), 0) + # Gumbel sampling needs logits, so we compute it in log space: + # log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0)) # The more numerically stable form is: - # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0)) + # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0)) ratio = tl.exp(draft_log_probs - target_log_probs) residual_logits = tl.where( ratio < 1.0, - target_log_probs + tl.log(1 - ratio), + target_log_probs + tldevice.log1p(-ratio), float("-inf"), ).to(tl.float32) else: From f9e684499f67071641bb2333d52dadd879231ac2 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 26 Jun 2026 07:59:57 +0800 Subject: [PATCH 008/138] [Rust Frontend] Migrate gemma4 to unified parser (#46602) Signed-off-by: Bugen Zhao --- pyproject.toml | 3 +- rust/Cargo.toml | 2 +- rust/src/chat/src/output/default/mod.rs | 164 ++++++- .../chat/src/output/default/structural_tag.rs | 24 +- rust/src/chat/src/output/default/unified.rs | 2 +- rust/src/chat/src/parser/mod.rs | 1 + rust/src/chat/src/parser/reasoning/mod.rs | 31 +- rust/src/chat/src/parser/reasoning/tests.rs | 2 + rust/src/chat/src/parser/tool/mod.rs | 29 +- rust/src/chat/src/parser/unified.rs | 124 ++++++ rust/src/chat/tests/roundtrip.rs | 53 ++- rust/src/parser/benches/gemma4.rs | 8 +- rust/src/parser/benches/utils/adapter.rs | 106 +++++ rust/src/parser/benches/utils/mod.rs | 7 + rust/src/parser/src/lib.rs | 1 + rust/src/parser/src/reasoning/delimited.rs | 10 +- rust/src/parser/src/reasoning/gemma4.rs | 273 ------------ rust/src/parser/src/reasoning/mod.rs | 8 +- rust/src/parser/src/tool/error.rs | 4 + rust/src/parser/src/tool/mod.rs | 4 +- .../parser/src/{tool => unified}/gemma4.rs | 409 ++++++++++++++++-- rust/src/parser/src/unified/mod.rs | 10 +- rust/src/parser/src/{tool => }/utils.rs | 119 ++++- 23 files changed, 1000 insertions(+), 394 deletions(-) create mode 100644 rust/src/chat/src/parser/unified.rs create mode 100644 rust/src/parser/benches/utils/adapter.rs delete mode 100644 rust/src/parser/src/reasoning/gemma4.rs rename rust/src/parser/src/{tool => unified}/gemma4.rs (64%) rename rust/src/parser/src/{tool => }/utils.rs (83%) diff --git a/pyproject.toml b/pyproject.toml index 249832ff2e5..3819ad7fc8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/parser/src/tool/gemma4.rs", "rust/src/text/src/output/decoded.rs", + "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs", + "rust/src/text/src/output/decoded.rs", "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] ignore-hidden = false diff --git a/rust/Cargo.toml b/rust/Cargo.toml index dc3895c372d..601e009df07 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -120,7 +120,7 @@ vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } -winnow = "1.0.2" +winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.1.0" zeromq = { version = "0.6.0", default-features = false, features = [ "tokio-runtime", diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 8c9a4362d1c..c494df600f4 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -18,7 +18,8 @@ use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEvent use crate::parser::ParserSelection; use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; -use crate::request::ChatRequest; +use crate::parser::unified::UnifiedParserFactory; +use crate::request::{ChatRequest, ChatTool}; use crate::{Error, Result as ChatResult}; /// Default request-scoped output processor used by Hugging Face style chat @@ -46,20 +47,31 @@ impl DefaultChatOutputProcessor { tool_call_parser: &ParserSelection, reasoning_parser: &ParserSelection, ) -> ChatResult { - let tool_parsing_enabled = request.tool_parsing_enabled(); - let tool_parser = if tool_parsing_enabled { - Some(Self::resolve_tool_parser( - request, + let parser = if tool_call_parser == reasoning_parser + && let Some(parser) = Self::resolve_optional_unified_parser( + &request.tools, model_id, + tokenizer.clone(), tool_call_parser, - )?) + )? { + parser } else { - None + let tool_parsing_enabled = request.tool_parsing_enabled(); + let tool_parser = if tool_parsing_enabled { + Some(Self::resolve_tool_parser( + &request.tools, + model_id, + tool_call_parser, + )?) + } else { + None + }; + let reasoning_parser = + Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?; + Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box }; - let reasoning_parser = - Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?; - let parser: Box = - Box::new(CombinedParser::new(reasoning_parser, tool_parser)); + + apply_structural_tag_constraint(request, parser.structural_tag_model())?; if parser.preserve_special_tokens() { request.decode_options.skip_special_tokens = false; @@ -84,7 +96,7 @@ impl DefaultChatOutputProcessor { } fn resolve_tool_parser( - request: &mut ChatRequest, + tools: &[ChatTool], model_id: &str, selection: &ParserSelection, ) -> ChatResult> { @@ -100,14 +112,36 @@ impl DefaultChatOutputProcessor { ParserSelection::Explicit(name) => name.as_str(), }; - let parser = factory.create(parser_name, &request.tools)?; - - apply_structural_tag_constraint(request, parser.as_ref())?; + let parser = factory.create(parser_name, tools)?; TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser")); Ok(parser) } + fn resolve_optional_unified_parser( + tools: &[ChatTool], + model_id: &str, + tokenizer: DynTokenizer, + selection: &ParserSelection, + ) -> ChatResult>> { + let factory = UnifiedParserFactory::global(); + let parser_name = match selection { + ParserSelection::Auto => factory.resolve_name_for_model(model_id), + ParserSelection::None => None, + ParserSelection::Explicit(name) if factory.contains(name) => Some(name.as_str()), + ParserSelection::Explicit(_) => None, + }; + + let Some(parser_name) = parser_name else { + return Ok(None); + }; + + let parser = factory.create(parser_name, tools, tokenizer)?; + + UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser")); + Ok(Some(parser)) + } + fn resolve_optional_reasoning_parser( model_id: &str, tokenizer: DynTokenizer, @@ -134,6 +168,7 @@ impl DefaultChatOutputProcessor { static TOOL_PARSER_LOG_ONCE: Once = Once::new(); static REASONING_PARSER_LOG_ONCE: Once = Once::new(); +static UNIFIED_PARSER_LOG_ONCE: Once = Once::new(); impl ChatOutputProcessor for DefaultChatOutputProcessor { /// Transforms a raw generate-output token stream into structured chat @@ -149,3 +184,102 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor { Ok(structured.boxed()) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::Tokenizer; + + use super::DefaultChatOutputProcessor; + use crate::Error; + use crate::parser::ParserSelection; + use crate::request::ChatRequest; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + "<|channel>" => Some(1), + "" => Some(2), + _ => None, + } + } + } + + fn tokenizer() -> Arc { + Arc::new(FakeTokenizer) + } + + #[test] + fn equal_explicit_gemma4_uses_unified_parser() { + let mut request = ChatRequest::for_test(); + let selection = ParserSelection::Explicit("gemma4".to_string()); + + DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &selection, + &selection, + ) + .unwrap(); + } + + #[test] + fn auto_auto_gemma4_model_uses_unified_parser() { + let mut request = ChatRequest::for_test(); + + DefaultChatOutputProcessor::new( + &mut request, + "google/gemma-4-27b-it", + tokenizer(), + &ParserSelection::Auto, + &ParserSelection::Auto, + ) + .unwrap(); + } + + #[test] + fn mixed_gemma4_selection_uses_split_dummy_error() { + let mut request = ChatRequest::for_test(); + let error = match DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &ParserSelection::Auto, + &ParserSelection::Explicit("gemma4".to_string()), + ) { + Ok(_) => panic!("expected mixed Gemma4 parser selection to fail"), + Err(error) => error, + }; + + let Error::ParserInitialization { error, .. } = error else { + panic!("expected parser initialization error"); + }; + assert_eq!( + error.to_string(), + "`gemma4` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + ); + } +} diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index bdc5d8e14c2..6ba2458ca8d 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -2,12 +2,12 @@ use thiserror_ext::AsReport; use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams}; +use vllm_parser::tool::StructuralTagModel; use xgrammar_structural_tag::{ FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam, build_structural_tag, }; -use crate::parser::tool::ToolParser; use crate::request::{ChatRequest, ChatToolChoice}; use crate::{Error, Result as ChatResult}; @@ -15,9 +15,9 @@ use crate::{Error, Result as ChatResult}; /// support and the request's tool choice. pub(super) fn apply_structural_tag_constraint( request: &mut ChatRequest, - parser: &dyn ToolParser, + model: Option, ) -> ChatResult<()> { - let Some(model) = parser.structural_tag_model() else { + let Some(model) = model else { return Ok(()); }; let Some(tool_choice) = structural_tag_tool_choice(request) else { @@ -77,7 +77,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option vllm_parser::reasoning::Result>; +type ReasoningParserCreator = Arc< + dyn Fn(DynTokenizer) -> vllm_parser::reasoning::Result> + Send + Sync, +>; /// Registry and model matcher for reasoning parsers. pub type ReasoningParserFactory = ParserFactory; @@ -58,7 +59,7 @@ impl ReasoningParserFactory { .register_parser::(names::DEEPSEEK_R1) .register_parser::(names::DEEPSEEK_V3) .register_parser::(names::DEEPSEEK_V4) - .register_parser::(names::GEMMA4) + .register_unified_dummy(names::GEMMA4) .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) @@ -109,7 +110,17 @@ impl ReasoningParserFactory { where T: ReasoningParser + 'static, { - self.register_creator(name, T::create) + self.register_creator(name, Arc::new(T::create)) + } + + /// Register one unified-only parser name in the split reasoning registry. + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self { + let name = name.to_string(); + let registered_name = name.clone(); + self.register_creator( + ®istered_name, + Arc::new(move |_| Err(ReasoningError::DummyUnifiedParser { name: name.clone() })), + ) } /// Construct a parser from an exact name. @@ -124,7 +135,7 @@ impl ReasoningParserFactory { available_names: self.list(), })?; - creator(tokenizer).map_err(|error| crate::Error::ParserInitialization { + creator.as_ref()(tokenizer).map_err(|error| crate::Error::ParserInitialization { kind: "reasoning", name: name.to_string(), error: error.into(), diff --git a/rust/src/chat/src/parser/reasoning/tests.rs b/rust/src/chat/src/parser/reasoning/tests.rs index 58d987770c6..e6255d14a00 100644 --- a/rust/src/chat/src/parser/reasoning/tests.rs +++ b/rust/src/chat/src/parser/reasoning/tests.rs @@ -35,11 +35,13 @@ fn factory_contains_and_lists_registered_parsers() { assert!(factory.contains(names::SEED_OSS)); assert!(factory.contains(names::STEP3P5)); assert!(factory.contains(names::MINIMAX_M3)); + assert!(factory.contains(names::GEMMA4)); assert!(factory.list().contains(&names::QWEN3.to_string())); assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string())); assert!(factory.list().contains(&names::SEED_OSS.to_string())); assert!(factory.list().contains(&names::STEP3P5.to_string())); assert!(factory.list().contains(&names::MINIMAX_M3.to_string())); + assert!(factory.list().contains(&names::GEMMA4.to_string())); } #[test] diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 9884d1aca2a..a156d670248 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -1,13 +1,13 @@ //! Tool parser registration and selection boundary for `vllm-chat`. -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; pub use vllm_parser::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, - Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, - HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, - MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, - Qwen3CoderToolParser, Qwen3XmlToolParser, ToolParser, ToolParserError, + Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser, + Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, + MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, + Qwen3XmlToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -40,7 +40,8 @@ pub mod names { } /// Constructor signature for one registered tool parser implementation. -type ToolParserCreator = fn(&[ChatTool]) -> vllm_parser::tool::Result>; +type ToolParserCreator = + Arc vllm_parser::tool::Result> + Send + Sync>; /// Registry and model matcher for tool parsers. pub type ToolParserFactory = ParserFactory; @@ -65,7 +66,7 @@ impl ToolParserFactory { .register_parser::(names::DEEPSEEK_V4) .register_parser::(names::GLM45) .register_parser::(names::GLM47) - .register_parser::(names::GEMMA4) + .register_unified_dummy(names::GEMMA4) .register_parser::(names::GRANITE4) .register_parser::(names::HERMES) .register_parser::(names::HY_V3) @@ -126,7 +127,17 @@ impl ToolParserFactory { where T: ToolParser + 'static, { - self.register_creator(name, T::create) + self.register_creator(name, Arc::new(T::create)) + } + + /// Register one unified-only parser name in the split tool registry. + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self { + let name = name.to_string(); + let registered_name = name.clone(); + self.register_creator( + ®istered_name, + Arc::new(move |_| Err(ToolParserError::DummyUnifiedParser { name: name.clone() })), + ) } /// Construct a parser from an exact name. @@ -137,7 +148,7 @@ impl ToolParserFactory { available_names: self.list(), })?; - creator(tools).map_err(|error| crate::Error::ParserInitialization { + creator.as_ref()(tools).map_err(|error| crate::Error::ParserInitialization { kind: "tool", name: name.to_string(), error: error.into(), diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs new file mode 100644 index 00000000000..6456cfda754 --- /dev/null +++ b/rust/src/chat/src/parser/unified.rs @@ -0,0 +1,124 @@ +//! Unified parser registration and selection boundary for `vllm-chat`. + +use std::sync::LazyLock; + +pub use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParser}; +use vllm_tokenizer::DynTokenizer; + +use crate::parser::ParserFactory; +use crate::request::ChatTool; + +/// Canonical public names for registered unified parsers. +pub mod names { + pub const GEMMA4: &str = "gemma4"; +} + +/// Constructor signature for one registered unified parser implementation. +type UnifiedParserCreator = + fn(&[ChatTool], DynTokenizer) -> vllm_parser::unified::Result>; + +/// Registry and model matcher for unified parsers. +pub type UnifiedParserFactory = ParserFactory; + +impl UnifiedParserFactory { + /// Get the global unified parser factory with built-in registrations and + /// model mappings. + pub fn global() -> &'static Self { + static INSTANCE: LazyLock = LazyLock::new(UnifiedParserFactory::new); + &INSTANCE + } + + /// Create the default registry with built-in parser names and model + /// mappings. + pub fn new() -> Self { + let mut factory = Self::default(); + + factory.register_parser::(names::GEMMA4); + + factory + .register_pattern("gemma-4", names::GEMMA4) + .register_pattern("gemma4", names::GEMMA4); + + factory + } + + /// Register one parser type that exposes a static `create()` constructor. + pub fn register_parser(&mut self, name: &str) -> &mut Self + where + T: UnifiedParser + 'static, + { + self.register_creator(name, T::create) + } + + /// Construct a parser from an exact name. + pub fn create( + &self, + name: &str, + tools: &[ChatTool], + tokenizer: DynTokenizer, + ) -> crate::Result> { + let creator = self.creator(name).ok_or_else(|| crate::Error::ParserUnavailableByName { + kind: "unified", + name: name.to_string(), + available_names: self.list(), + })?; + + creator(tools, tokenizer).map_err(|error| crate::Error::ParserInitialization { + kind: "unified", + name: name.to_string(), + error: error.into(), + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::Tokenizer; + + use super::{UnifiedParserFactory, names}; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + "<|channel>" => Some(1), + "" => Some(2), + _ => None, + } + } + } + + #[test] + fn factory_registers_gemma4() { + let factory = UnifiedParserFactory::new(); + + assert!(factory.contains(names::GEMMA4)); + assert_eq!( + factory.resolve_name_for_model("google/gemma-4-27b-it"), + Some(names::GEMMA4) + ); + factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap(); + } +} diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index b3d5d9eae34..15bd4aca23a 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -37,6 +37,8 @@ struct RoundtripCase { /// JSON formatting expected after this model's template has materialized /// tool-call arguments. json_fmt: JsonFmt, + /// Whether the template renders tool-call argument object keys in sorted order. + sort_json_keys: bool, } #[derive(Clone, Copy)] @@ -81,6 +83,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), + sort_json_keys: false, } } @@ -93,6 +96,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -105,6 +109,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -117,6 +122,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: false }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -129,6 +135,20 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + + /// Gemma4 channel reasoning with custom function-call arguments. + fn gemma4() -> Self { + Self { + model_id: "google/gemma-4-E4B-it", + assistant_stop_suffix: "<|tool_response>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Always { value: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: true, } } @@ -142,6 +162,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), + sort_json_keys: false, } } @@ -154,6 +175,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -166,6 +188,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } } @@ -195,8 +218,8 @@ roundtrip_tests! { seed_oss => [reasoning_and_content], step3p5 => [reasoning_and_content], - // Note: Kimi K2.5 strips the reasoning content in history. - kimi_k25 => [tool_call_mix], + gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call + kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history } /// Run the fixed reasoning+content fixture for one model/parser case. @@ -347,14 +370,38 @@ fn spaced_json_fmt() -> JsonFmt { /// Pass in a raw JSON string instead of a structured value to ensure the exact precision and /// formatting of numbers are preserved. fn expected_arguments(case: &RoundtripCase, raw_json: &str) -> Result { - let value: serde_json::Value = + let mut value: serde_json::Value = serde_json::from_str(raw_json).context("invalid expected tool-call arguments")?; + if case.sort_json_keys { + sort_json_value(&mut value); + } case.json_fmt .format_to_string(&value) .context("failed to format expected tool-call arguments") } +/// Sort JSON object keys recursively to match templates that render mappings with `dictsort`. +fn sort_json_value(value: &mut serde_json::Value) { + match value { + serde_json::Value::Object(map) => { + for value in map.values_mut() { + sort_json_value(value); + } + + let mut entries = std::mem::take(map).into_iter().collect::>(); + entries.sort_by(|(left, _), (right, _)| left.cmp(right)); + map.extend(entries); + } + serde_json::Value::Array(values) => { + for value in values { + sort_json_value(value); + } + } + _ => {} + } +} + /// Load the real model chat/text backend for one roundtrip case. async fn load_roundtrip_backends(case: &RoundtripCase) -> Result { load_model_backends( diff --git a/rust/src/parser/benches/gemma4.rs b/rust/src/parser/benches/gemma4.rs index 761f8d4e235..fd29e77a9a2 100644 --- a/rust/src/parser/benches/gemma4.rs +++ b/rust/src/parser/benches/gemma4.rs @@ -2,10 +2,11 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; -use vllm_parser::tool::{Gemma4ToolParser, Tool, ToolParser}; +use vllm_parser::tool::{Tool, ToolParser}; +use vllm_parser::unified::Gemma4UnifiedParser; mod utils; -use utils::feed_parser; +use utils::{UnifiedToolParserAdapter, feed_parser}; const CHUNK_CHARS: usize = 7; const LONG_NORMAL_TEXT_REPEATS: usize = 2048; @@ -68,7 +69,8 @@ fn long_tool_argument_fixture() -> String { } fn parser(tools: &[Tool]) -> Box { - Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize") + UnifiedToolParserAdapter::::create(tools) + .expect("Gemma4 unified parser should initialize") } fn run_stream_group( diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs new file mode 100644 index 00000000000..103ed18d093 --- /dev/null +++ b/rust/src/parser/benches/utils/adapter.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +use vllm_parser::tool::{ + Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput, +}; +use vllm_parser::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, +}; +use vllm_tokenizer::Tokenizer; + +/// Tokenizer stub used by unified-parser benchmarks. +struct BenchTokenizer; + +impl Tokenizer for BenchTokenizer { + fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result> { + Ok(text.chars().map(|_| u32::MAX).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok("\u{FFFD}".repeat(token_ids.len())) + } + + fn token_to_id(&self, _token: &str) -> Option { + Some(u32::MAX) + } +} + +/// Bench-only adapter that exposes a unified parser through the tool-parser +/// benchmark harness. +/// +/// Returns error if the unified parser produces reasoning events. +pub struct UnifiedToolParserAdapter { + inner: Box, + _marker: std::marker::PhantomData, +} + +fn map_unified_error(error: UnifiedParserError) -> ToolParserError { + ToolParserError::ParsingFailed { + message: format!("unified parser failed: {error}"), + } +} + +fn append_unified_output( + output: UnifiedParserOutput, + tool_output: &mut ToolParserOutput, +) -> Result<()> { + for event in output.events { + match event { + UnifiedParserEvent::Text(text) => tool_output.push_text(text), + UnifiedParserEvent::ToolCall(call) => tool_output.push_call(call), + UnifiedParserEvent::Reasoning(_) => { + return Err(ToolParserError::ParsingFailed { + message: "unified parser emitted reasoning in tool-parser adapter".to_string(), + }); + } + } + } + Ok(()) +} + +impl ToolParser for UnifiedToolParserAdapter { + fn create(tools: &[Tool]) -> Result> + where + Self: Sized + 'static, + { + let inner = T::create(tools, Arc::new(BenchTokenizer)).map_err(map_unified_error)?; + Ok(Box::new(Self { + inner, + _marker: std::marker::PhantomData, + })) + } + + fn preserve_special_tokens(&self) -> bool { + self.inner.preserve_special_tokens() + } + + fn structural_tag_model(&self) -> Option { + self.inner.structural_tag_model() + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.inner.tool_call_id(tool_index) + } + + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + let mut unified_output = UnifiedParserOutput::default(); + let result = self.inner.parse_into(chunk, &mut unified_output).map_err(map_unified_error); + append_unified_output(unified_output, output)?; + result + } + + fn finish(&mut self) -> Result { + let unified_output = self.inner.finish().map_err(map_unified_error)?; + let mut output = ToolParserOutput::default(); + append_unified_output(unified_output, &mut output)?; + Ok(output) + } + + fn reset(&mut self) -> String { + self.inner.reset() + } +} diff --git a/rust/src/parser/benches/utils/mod.rs b/rust/src/parser/benches/utils/mod.rs index 914766a79aa..229f40a3681 100644 --- a/rust/src/parser/benches/utils/mod.rs +++ b/rust/src/parser/benches/utils/mod.rs @@ -1,4 +1,9 @@ +// This module is shared by multiple benchmark targets. +// There could be false positives for unused code or imports, and fixing them would lead to some other benchmarks failing to compile. #![allow(dead_code)] +#![allow(unused_imports)] + +mod adapter; use futures::FutureExt as _; use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool}; @@ -6,6 +11,8 @@ use tool_parser::traits::ToolParser as ExternalToolParser; use vllm_parser::tool::test_utils::collect_stream; use vllm_parser::tool::{Tool, ToolParser}; +pub(super) use adapter::UnifiedToolParserAdapter; + pub(super) fn openai_tools(tools: &[Tool]) -> Vec { tools .iter() diff --git a/rust/src/parser/src/lib.rs b/rust/src/parser/src/lib.rs index 5ba2cf60edd..0b5c2b6d782 100644 --- a/rust/src/parser/src/lib.rs +++ b/rust/src/parser/src/lib.rs @@ -3,3 +3,4 @@ pub mod reasoning; pub mod tool; pub mod unified; +pub(crate) mod utils; diff --git a/rust/src/parser/src/reasoning/delimited.rs b/rust/src/parser/src/reasoning/delimited.rs index 69b4db5f183..256e95fdde3 100644 --- a/rust/src/parser/src/reasoning/delimited.rs +++ b/rust/src/parser/src/reasoning/delimited.rs @@ -144,20 +144,20 @@ impl DelimitedReasoningParser { } /// Determine the reasoning state implied by the last prompt boundary, if any. -fn last_reasoning_boundary( +pub(crate) fn last_reasoning_boundary( prompt_token_ids: &[u32], start_token_id: u32, end_token_id: u32, tokenizer: &dyn Tokenizer, ) -> Option { - for token_id in prompt_token_ids.iter().rev() { - if *token_id == start_token_id { + for token_id in prompt_token_ids.iter().rev().copied() { + if token_id == start_token_id { return Some(true); } - if *token_id == end_token_id { + if token_id == end_token_id { return Some(false); } - if tokenizer.is_special_id(*token_id) { + if tokenizer.is_special_id(token_id) { return None; } } diff --git a/rust/src/parser/src/reasoning/gemma4.rs b/rust/src/parser/src/reasoning/gemma4.rs deleted file mode 100644 index ac5a6a17165..00000000000 --- a/rust/src/parser/src/reasoning/gemma4.rs +++ /dev/null @@ -1,273 +0,0 @@ -use vllm_tokenizer::DynTokenizer; - -use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; - -const THOUGHT_PREFIX: &str = "thought\n"; - -/// Reasoning parser for Google Gemma4 thinking models. -/// -/// Gemma4 emits reasoning inside `<|channel> ... ` spans and adds a -/// structural `thought\n` label at the beginning of the reasoning channel. -/// This parser keeps the delimiter handling in the shared delimited parser and -/// only layers on Gemma4-specific request adjustment plus prefix stripping. -/// -/// Original Python implementation: -/// -pub struct Gemma4ReasoningParser { - inner: DelimitedReasoningParser, - reasoning_text: String, - prefix_stripped: bool, -} - -impl Gemma4ReasoningParser { - /// Create a Gemma4 parser. - pub fn new(tokenizer: DynTokenizer) -> Result { - Ok(Self { - inner: DelimitedReasoningParser::new(tokenizer, "<|channel>", "", false)?, - reasoning_text: String::new(), - prefix_stripped: false, - }) - } - - /// Apply Gemma4's `thought\n` stripping rule to one reasoning delta. - /// - /// Early reasoning text is buffered until we can decide whether it begins - /// with the structural channel label. - fn strip_thought_prefix(&mut self, reasoning: &str) -> Option { - if self.prefix_stripped { - return Some(reasoning.to_string()); - } - - self.reasoning_text.push_str(reasoning); - - if self.reasoning_text.starts_with(THOUGHT_PREFIX) { - let prefix_len = THOUGHT_PREFIX.len(); - let previous_len = self.reasoning_text.len() - reasoning.len(); - if previous_len >= prefix_len { - self.reasoning_text.clear(); - self.prefix_stripped = true; - return Some(reasoning.to_string()); - } - - let prefix_chars_in_delta = prefix_len - previous_len; - let stripped = &reasoning[prefix_chars_in_delta.min(reasoning.len())..]; - if stripped.is_empty() { - if self.reasoning_text.len() >= prefix_len { - self.reasoning_text.clear(); - self.prefix_stripped = true; - } - return None; - } - - self.reasoning_text.clear(); - self.prefix_stripped = true; - return Some(stripped.to_string()); - } - - if THOUGHT_PREFIX.starts_with(&self.reasoning_text) { - return None; - } - - self.prefix_stripped = true; - Some(std::mem::take(&mut self.reasoning_text)) - } - - /// Apply Gemma4-specific reasoning post-processing to one parsed delta. - fn post_process(&mut self, mut result: ReasoningDelta) -> ReasoningDelta { - if let Some(reasoning) = result.reasoning.take() { - result.reasoning = - self.strip_thought_prefix(&reasoning).filter(|text| !text.is_empty()); - } - result - } -} - -impl ReasoningParser for Gemma4ReasoningParser { - fn create(tokenizer: DynTokenizer) -> Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self::new(tokenizer)?)) - } - - fn preserve_special_tokens(&self) -> bool { - true - } - - fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { - self.inner.initialize(prompt_token_ids); - self.reasoning_text.clear(); - self.prefix_stripped = false; - Ok(()) - } - - fn push(&mut self, delta: &str) -> Result { - let result = self.inner.push(delta); - Ok(self.post_process(result)) - } - - fn finish(&mut self) -> Result { - let result = self.inner.finish(); - Ok(self.post_process(result)) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use vllm_tokenizer::Tokenizer; - - use super::Gemma4ReasoningParser; - use crate::reasoning::ReasoningParser; - - struct FakeTokenizer; - - impl Tokenizer for FakeTokenizer { - fn encode( - &self, - text: &str, - _add_special_tokens: bool, - ) -> vllm_tokenizer::Result> { - Ok(text.chars().map(u32::from).collect()) - } - - fn decode( - &self, - token_ids: &[u32], - _skip_special_tokens: bool, - ) -> vllm_tokenizer::Result { - Ok(token_ids - .iter() - .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) - .collect()) - } - - fn token_to_id(&self, token: &str) -> Option { - match token { - "<|channel>" => Some(1000), - "" => Some(1001), - _ => None, - } - } - } - - fn run_streaming(output: &[&str]) -> (Option, Option) { - let tokenizer = Arc::new(FakeTokenizer); - let mut parser = Gemma4ReasoningParser::new(tokenizer).unwrap(); - let mut reasoning = String::new(); - let mut content = String::new(); - - for delta in output { - let result = parser.push(delta).unwrap(); - if let Some(next) = result.reasoning { - reasoning.push_str(&next); - } - if let Some(next) = result.content { - content.push_str(&next); - } - } - - let final_delta = parser.finish().unwrap(); - if let Some(next) = final_delta.reasoning { - reasoning.push_str(&next); - } - if let Some(next) = final_delta.content { - content.push_str(&next); - } - - ( - (!reasoning.is_empty()).then_some(reasoning), - (!content.is_empty()).then_some(content), - ) - } - - #[test] - fn gemma4_reasoning_streaming_handles_channel_delimited_outputs() { - let cases = [ - ( - "no_reasoning", - vec!["This is content"], - None, - Some("This is content"), - ), - ( - "reasoning_and_content", - vec!["<|channel>This is a reasoning sectionThis is the rest"], - Some("This is a reasoning section"), - Some("This is the rest"), - ), - ( - "complete_reasoning", - vec!["<|channel>This is a reasoning section"], - Some("This is a reasoning section"), - None, - ), - ( - "multiple_lines", - vec!["<|channel>This\nThatThis is the rest\nThat"], - Some("This\nThat"), - Some("This is the rest\nThat"), - ), - ( - "no_end", - vec!["<|channel>This is a reasoning section"], - Some("This is a reasoning section"), - None, - ), - ("empty", vec![""], None, None), - ( - "newline_around_reasoning", - vec!["Before\n<|channel>This is a reasoning section\nThis is the rest"], - Some("This is a reasoning section"), - Some("Before\n\nThis is the rest"), - ), - ( - "thought_prefix", - vec!["<|channel>thought\nActual reasoning hereFinal answer"], - Some("Actual reasoning here"), - Some("Final answer"), - ), - ( - "thought_prefix_only", - vec!["<|channel>thought\n"], - None, - None, - ), - ( - "thought_prefix_multiline", - vec!["<|channel>thought\nLine1\nLine2Answer"], - Some("Line1\nLine2"), - Some("Answer"), - ), - ( - "thought_prefix_diverge", - vec!["<|channel>thousand reasonsDone"], - Some("thousand reasons"), - Some("Done"), - ), - ]; - - for (name, output, expected_reasoning, expected_content) in cases { - let (reasoning, content) = run_streaming(&output); - assert_eq!(reasoning.as_deref(), expected_reasoning, "{name}"); - assert_eq!(content.as_deref(), expected_content, "{name}"); - } - } - - #[test] - fn gemma4_strips_thought_prefix_even_when_split_across_deltas() { - let (reasoning, content) = - run_streaming(&["<|channel>thou", "ght", "\nabc", "done"]); - assert_eq!(reasoning.as_deref(), Some("abc")); - assert_eq!(content.as_deref(), Some("done")); - } - - #[test] - fn gemma4_preserves_special_tokens() { - let tokenizer = Arc::new(FakeTokenizer); - let parser = Gemma4ReasoningParser::new(tokenizer).unwrap(); - - assert!(parser.preserve_special_tokens()); - } -} diff --git a/rust/src/parser/src/reasoning/mod.rs b/rust/src/parser/src/reasoning/mod.rs index 1f71e14cef7..fcb0f96792a 100644 --- a/rust/src/parser/src/reasoning/mod.rs +++ b/rust/src/parser/src/reasoning/mod.rs @@ -17,7 +17,6 @@ mod cohere_cmd; mod deepseek_r1; mod delimited; -mod gemma4; mod kimi; mod minimax_m3; mod qwen3; @@ -29,8 +28,7 @@ use vllm_tokenizer::DynTokenizer; pub use self::cohere_cmd::CohereCmdReasoningParser; pub use self::deepseek_r1::DeepSeekR1ReasoningParser; -pub(crate) use self::delimited::DelimitedReasoningParser; -pub use self::gemma4::Gemma4ReasoningParser; +pub(crate) use self::delimited::{DelimitedReasoningParser, last_reasoning_boundary}; pub use self::kimi::KimiReasoningParser; pub use self::minimax_m3::MiniMaxM3ReasoningParser; pub use self::qwen3::Qwen3ReasoningParser; @@ -129,6 +127,10 @@ pub trait ReasoningParser: Send { pub enum ReasoningError { #[error("tokenizer is missing reasoning delimiter token `{token}`")] MissingToken { token: String }, + #[error( + "`{name}` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + )] + DummyUnifiedParser { name: String }, } #[cfg(test)] diff --git a/rust/src/parser/src/tool/error.rs b/rust/src/parser/src/tool/error.rs index 6a64c257d8c..4b2b4efcb45 100644 --- a/rust/src/parser/src/tool/error.rs +++ b/rust/src/parser/src/tool/error.rs @@ -10,4 +10,8 @@ pub type Result = std::result::Result; pub enum ToolParserError { #[error("tool parser parsing failed: {message}")] ParsingFailed { message: String }, + #[error( + "`{name}` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + )] + DummyUnifiedParser { name: String }, } diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index a27e202e660..dd4630b1c6b 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -4,7 +4,6 @@ pub(crate) mod error; mod deepseek_dsml; pub(crate) mod deepseek_json; -mod gemma4; mod glm_xml; mod hy_v3; mod json; @@ -15,14 +14,13 @@ mod parameters; mod qwen_coder; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; -pub(crate) mod utils; +use crate::utils; use std::collections::{BTreeMap, btree_map}; pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser}; pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser}; pub use error::{Result, ToolParserError}; -pub use gemma4::Gemma4ToolParser; pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser}; pub use hy_v3::HyV3ToolParser; pub use json::{ diff --git a/rust/src/parser/src/tool/gemma4.rs b/rust/src/parser/src/unified/gemma4.rs similarity index 64% rename from rust/src/parser/src/tool/gemma4.rs rename to rust/src/parser/src/unified/gemma4.rs index e5a95485ce8..3276088913d 100644 --- a/rust/src/parser/src/tool/gemma4.rs +++ b/rust/src/parser/src/unified/gemma4.rs @@ -6,10 +6,17 @@ use winnow::prelude::*; use winnow::stream::{Partial, Stream}; use winnow::token::{literal, take_till, take_until}; -use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len}; -use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::tool::Tool; +use vllm_tokenizer::DynTokenizer; +use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; +use crate::reasoning::last_reasoning_boundary; +use crate::tool::{Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len_mul}; + +const REASONING_START: &str = "<|channel>thought\n"; +const CHANNEL_START: &str = "<|channel>"; +const CHANNEL_END: &str = ""; const TOOL_CALL_START: &str = "<|tool_call>"; const TOOL_CALL_END: &str = ""; const STRING_DELIM: &str = "<|\"|>"; @@ -20,6 +27,9 @@ type Gemma4Input<'i> = Partial<&'i str>; #[derive(Debug, Clone, PartialEq)] enum Gemma4Event { Text { len: usize }, + Reasoning { len: usize }, + ReasoningStart, + ReasoningEnd, ToolCallStart, ToolCallHeader { name: String }, ToolCall { args: Map }, @@ -35,6 +45,7 @@ struct Gemma4ArgsScanState { enum Gemma4Mode { #[default] Text, + Reasoning, Header, ToolCall { name: String, @@ -42,36 +53,62 @@ enum Gemma4Mode { }, } -/// Tool parser for Google Gemma4 models. +/// Unified parser for Google Gemma4 models. /// /// Original Python implementation: -/// +/// /// -/// Handles the Gemma4 function call format: +/// Handles Gemma4 reasoning and function-call formats: +/// +/// `<|channel>thought\nreasoning` /// /// `<|tool_call>call:func_name{key:<|"|>value<|"|>}` /// /// Arguments are emitted only after a full Gemma4 tool call is parsed. -pub struct Gemma4ToolParser { +pub struct Gemma4UnifiedParser { buffer: String, mode: Gemma4Mode, emitted_tool_count: usize, + tokenizer: DynTokenizer, + channel_start_token_id: u32, + channel_end_token_id: u32, } -impl Gemma4ToolParser { - fn new(_tools: &[Tool]) -> Self { - Self { +impl Gemma4UnifiedParser { + /// Create a Gemma4 parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let channel_start_token_id = tokenizer.token_to_id(CHANNEL_START).ok_or_else(|| { + UnifiedParserError::MissingToken { + token: CHANNEL_START.to_string(), + } + })?; + let channel_end_token_id = + tokenizer + .token_to_id(CHANNEL_END) + .ok_or_else(|| UnifiedParserError::MissingToken { + token: CHANNEL_END.to_string(), + })?; + + Ok(Self { buffer: String::new(), mode: Gemma4Mode::default(), emitted_tool_count: 0, - } + channel_start_token_id, + channel_end_token_id, + tokenizer, + }) } - fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> { + fn apply_event(&mut self, event: Gemma4Event, output: &mut UnifiedParserOutput) -> Result<()> { match event { Gemma4Event::Text { len: consumed_len } => { - output.push_text(&self.buffer[..consumed_len]); + output.push_text(self.buffer[..consumed_len].to_string()); } + Gemma4Event::Reasoning { len: consumed_len } => { + output.push_reasoning(self.buffer[..consumed_len].to_string()); + } + Gemma4Event::ReasoningStart => self.mode = Gemma4Mode::Reasoning, + Gemma4Event::ReasoningEnd => self.mode = Gemma4Mode::Text, Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header, Gemma4Event::ToolCallHeader { name } => { self.mode = Gemma4Mode::ToolCall { @@ -100,9 +137,24 @@ impl Gemma4ToolParser { Ok(()) } + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = match last_reasoning_boundary( + prompt_token_ids, + self.channel_start_token_id, + self.channel_end_token_id, + self.tokenizer.as_ref(), + ) { + Some(true) => Gemma4Mode::Reasoning, + Some(false) | None => Gemma4Mode::Text, + }; + } + fn reset(&mut self) -> String { let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) { Gemma4Mode::Text => std::mem::take(&mut self.buffer), + Gemma4Mode::Reasoning => { + format!("{}{}", REASONING_START, std::mem::take(&mut self.buffer)) + } Gemma4Mode::Header => { format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer)) } @@ -122,19 +174,26 @@ impl Gemma4ToolParser { } } -impl ToolParser for Gemma4ToolParser { - fn create(tools: &[Tool]) -> Result> +impl UnifiedParser for Gemma4UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> where Self: Sized + 'static, { - Ok(Box::new(Self::new(tools))) + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.emitted_tool_count = 0; + self.initialize_mode(prompt_token_ids); + Ok(()) } fn preserve_special_tokens(&self) -> bool { true } - fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { self.buffer.push_str(chunk); while let Some((event, consumed_len)) = { @@ -149,11 +208,12 @@ impl ToolParser for Gemma4ToolParser { Ok(()) } - fn finish(&mut self) -> Result { - let mut output = ToolParserOutput::default(); + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); match &self.mode { - Gemma4Mode::Text => output.push_text(&self.buffer), + Gemma4Mode::Text => output.push_text(std::mem::take(&mut self.buffer)), + Gemma4Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => { return Err(parsing_failed!("incomplete Gemma4 tool call")); } @@ -164,7 +224,7 @@ impl ToolParser for Gemma4ToolParser { } fn reset(&mut self) -> String { - Gemma4ToolParser::reset(self) + Gemma4UnifiedParser::reset(self) } } @@ -175,6 +235,7 @@ fn parse_next_gemma4_event( ) -> ModalResult { match mode { Gemma4Mode::Text => parse_text_event(input), + Gemma4Mode::Reasoning => parse_reasoning_event(input), Gemma4Mode::Header => tool_call_header_event(input), Gemma4Mode::ToolCall { args_scan, .. } => tool_call_args_event(input, args_scan), } @@ -182,7 +243,32 @@ fn parse_next_gemma4_event( /// Parse a Gemma4 text-mode event. fn parse_text_event(input: &mut Gemma4Input<'_>) -> ModalResult { - alt((tool_call_start_event, safe_text_event)).parse_next(input) + alt(( + reasoning_start_event, + tool_call_start_event, + safe_text_event, + )) + .parse_next(input) +} + +/// Parse a Gemma4 reasoning-mode event. +fn parse_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult { + alt(( + reasoning_end_event, + tool_call_start_event, + safe_reasoning_event, + )) + .parse_next(input) +} + +/// Parse a Gemma4 reasoning start marker. +fn reasoning_start_event(input: &mut Gemma4Input<'_>) -> ModalResult { + literal(REASONING_START).value(Gemma4Event::ReasoningStart).parse_next(input) +} + +/// Parse a Gemma4 reasoning end marker. +fn reasoning_end_event(input: &mut Gemma4Input<'_>) -> ModalResult { + literal(CHANNEL_END).value(Gemma4Event::ReasoningEnd).parse_next(input) } /// Parse a Gemma4 tool-call start marker. @@ -226,7 +312,14 @@ fn gemma4_tool_name(input: &mut Gemma4Input<'_>) -> ModalResult { /// Parse a safe text run before the next Gemma4 marker. fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult { - safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len }) + safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START]) + .map(|len| Gemma4Event::Text { len }) +} + +/// Parse a safe reasoning run before the next Gemma4 marker. +fn safe_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult { + safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START]) + .map(|len| Gemma4Event::Reasoning { len }) } /// Parse raw Gemma4 arguments through the first end marker outside a Gemma string. @@ -418,17 +511,145 @@ fn parse_gemma4_scalar(value: &str) -> Value { #[cfg(test)] mod tests { + use std::sync::Arc; + use serde_json::{Value, json}; use thiserror_ext::AsReport; + use vllm_tokenizer::Tokenizer; use winnow::combinator::{eof, terminated}; use winnow::error::ErrMode; use winnow::prelude::*; use super::{ - Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content, - parse_gemma4_args, + CHANNEL_END, CHANNEL_START, Gemma4UnifiedParser, ToolCallDelta, UnifiedParser, + UnifiedParserError, UnifiedParserOutput, gemma4_array_content, parse_gemma4_args, }; - use crate::tool::{Tool, ToolParserTestExt as _}; + use crate::tool::Tool; + use crate::unified::{UnifiedParserEvent, parsing_failed}; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + CHANNEL_START => Some(100), + CHANNEL_END => Some(101), + _ => None, + } + } + + fn is_special_id(&self, token_id: u32) -> bool { + matches!(token_id, 100..=105) + } + } + + struct MissingTokenTokenizer; + + impl Tokenizer for MissingTokenTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, _token: &str) -> Option { + None + } + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for Gemma4UnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec<&ToolCallDelta>; + fn coalesce(self) -> Self; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + UnifiedParserEvent::Reasoning(_) | UnifiedParserEvent::ToolCall(_) => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + UnifiedParserEvent::Text(_) | UnifiedParserEvent::ToolCall(_) => None, + }) + .collect() + } + + fn calls(&self) -> Vec<&ToolCallDelta> { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(_) | UnifiedParserEvent::Reasoning(_) => None, + UnifiedParserEvent::ToolCall(call) => Some(call), + }) + .collect() + } + + fn coalesce(self) -> Self { + self + } + } fn parse_gemma4_array(array: &str) -> super::Result> { let mut input = array; @@ -494,9 +715,26 @@ mod tests { ] } - fn collect_stream(chunks: &[&str]) -> ToolParserOutput { - let mut parser = Gemma4ToolParser::new(&test_tools()); - let mut output = ToolParserOutput::default(); + fn test_parser() -> Gemma4UnifiedParser { + Gemma4UnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap() + } + + #[test] + fn gemma4_create_requires_channel_start_token() { + let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(MissingTokenTokenizer)) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == CHANNEL_START + )); + } + + fn collect_stream(chunks: &[&str]) -> UnifiedParserOutput { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); for chunk in chunks { output.append(parser.parse_chunk(chunk).unwrap()); } @@ -504,7 +742,7 @@ mod tests { output.coalesce() } - fn first_call(output: &ToolParserOutput) -> ToolCallDelta { + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { (*output.calls().first().expect("expected one tool call")).clone() } @@ -542,7 +780,7 @@ mod tests { #[test] fn gemma4_parse_complete_extracts_single_tool_call() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let output = parser .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}") .unwrap(); @@ -558,7 +796,7 @@ mod tests { #[test] fn gemma4_parse_complete_rejects_incomplete_tool_call() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let error = parser .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London") .unwrap_err(); @@ -607,8 +845,8 @@ mod tests { #[test] fn gemma4_streaming_waits_for_complete_tool_call() { - let mut parser = Gemma4ToolParser::new(&test_tools()); - let mut output = ToolParserOutput::default(); + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); for chunk in [ "<|tool_call>", @@ -773,7 +1011,7 @@ mod tests { #[test] fn gemma4_finish_flushes_partial_start_marker_as_text() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let mut output = parser.parse_chunk("<").unwrap(); output.append(parser.finish().unwrap()); @@ -781,9 +1019,104 @@ mod tests { assert!(output.calls().is_empty()); } + #[test] + fn gemma4_streaming_emits_reasoning_then_text() { + let output = collect_stream(&["<|channel>thought\nreasonanswer"]); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + assert!(output.calls().is_empty()); + } + + #[test] + fn gemma4_streaming_holds_split_reasoning_start() { + let mut parser = test_parser(); + + let first = parser.parse_chunk("<|channel>").unwrap(); + assert!(first.events.is_empty()); + + let mut output = parser.parse_chunk("thought\nrea").unwrap(); + output.append(parser.parse_chunk("sonanswer").unwrap()); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() { + let mut parser = test_parser(); + parser.initialize(&[100, 3000, 3001]).unwrap(); + + let output = parser.parse_complete("reasonanswer").unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_turn_prompt_starts_in_text() { + let mut parser = test_parser(); + parser.initialize(&[104, 3000, 3001]).unwrap(); + + let output = parser.parse_complete("<|channel>thought\nreasonanswer").unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_special_token_caps_boundary_scan() { + let mut parser = test_parser(); + parser.initialize(&[100, 3000, 104, 3001]).unwrap(); + + let output = parser.parse_complete("answer").unwrap(); + + assert!(output.reasoning_text().is_empty()); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_closed_channel_prompt_starts_in_text() { + let mut parser = test_parser(); + parser.initialize(&[100, 3000, 3001, 101]).unwrap(); + + let output = parser.parse_complete("answer").unwrap(); + + assert!(output.reasoning_text().is_empty()); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_reasoning_tool_call_implicitly_ends_reasoning() { + let output = collect_stream(&[ + "<|channel>thought\nNeed weather.", + "<|tool_call>", + "call:get_weather{location:<|\"|>Paris<|\"|>}", + "", + ]); + + assert_eq!(output.reasoning_text(), "Need weather."); + assert!(output.normal_text().is_empty()); + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "location": "Paris" }) + ); + } + + #[test] + fn gemma4_bare_channel_start_is_plain_text() { + let output = collect_stream(&["<|channel>plain"]); + + assert_eq!(output.normal_text(), "<|channel>plain"); + assert!(output.reasoning_text().is_empty()); + assert!(output.calls().is_empty()); + } + #[test] fn gemma4_finish_rejects_complete_args_without_end_marker() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); for chunk in ["<|tool_call>", "call:get_status{}"] { parser.parse_chunk(chunk).unwrap(); } @@ -795,7 +1128,7 @@ mod tests { #[test] fn gemma4_reset_preserves_internally_buffered_arguments() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); for chunk in [ "<|tool_call>", "call:write_file{", @@ -815,7 +1148,7 @@ mod tests { #[test] fn gemma4_reset_preserves_completed_arguments_after_parse_error() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let input = "<|tool_call>call:set{broken}"; let _error = parser.parse_chunk(input).unwrap_err(); diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 49955b4d818..6fe7d29b879 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -1,11 +1,14 @@ //! Unified parser interface for reasoning and tool-call deltas. mod combined; +mod gemma4; use thiserror::Error; +use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; pub use combined::CombinedParser; +pub use gemma4::Gemma4UnifiedParser; use crate::reasoning::ReasoningError; use crate::tool::{ @@ -189,10 +192,15 @@ pub trait UnifiedParser: Send { } /// Errors produced while creating or running unified parsers. -#[derive(Debug, Error)] +#[derive(Debug, Error, Macro)] +#[thiserror_ext(macro(path = "crate::unified", mangle))] pub enum UnifiedParserError { #[error("combined parser is constructed from split parser instances")] CombinedParserConstructor, + #[error("tokenizer is missing unified parser token `{token}`")] + MissingToken { token: String }, + #[error("unified parser parsing failed: {message}")] + ParsingFailed { message: String }, #[error(transparent)] Reasoning(#[from] ReasoningError), #[error(transparent)] diff --git a/rust/src/parser/src/tool/utils.rs b/rust/src/parser/src/utils.rs similarity index 83% rename from rust/src/parser/src/tool/utils.rs rename to rust/src/parser/src/utils.rs index b545c5881c1..70255215393 100644 --- a/rust/src/parser/src/tool/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -1,10 +1,10 @@ -//! Shared helpers for tool parsers. +//! Shared helpers for streaming parsers. use winnow::Parser; use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue}; -use winnow::stream::{Offset, Partial, Stream}; +use winnow::stream::{FindSlice, Offset, Partial, Stream}; -use super::Result; +use crate::tool::{Result, ToolParserError}; /// Return the byte length of the longest proper prefix of `token` that is also /// a suffix of `buffer`. @@ -15,7 +15,7 @@ use super::Result; /// The returned length is always a valid UTF-8 boundary in `token`, so callers /// can safely slice `&token[..len]` even when markers contain non-ASCII /// characters such as DeepSeek's DSML delimiters. -pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize { +pub fn partial_prefix_len(buffer: &str, token: &str) -> usize { let Some(first_byte) = token.as_bytes().first().copied() else { return 0; }; @@ -44,9 +44,10 @@ pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize { } /// Parse a safe text run before the next marker. +/// This is the single-marker variant of [`safe_text_len_mul`]. /// /// Returns the text length in bytes, and advances the input. -pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult { +pub fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult { let text = **input; if text.is_empty() { return incomplete(); @@ -67,15 +68,53 @@ pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalRes Ok(emit_len) } +/// Parse a safe text run before the earliest next marker. +/// This is the multi-marker variant of [`safe_text_len`]. +/// +/// Returns the text length in bytes, and advances the input. +pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalResult { + let text = **input; + if text.is_empty() { + return incomplete(); + } + + if let Some(start_idx) = find_slice_mul(text, markers) { + input.next_slice(start_idx); + return Ok(start_idx); + } + + let keep_len = markers.iter().map(|marker| partial_prefix_len(text, marker)).max().unwrap_or(0); + let emit_len = text.len().saturating_sub(keep_len); + if emit_len == 0 { + return incomplete(); + } + + input.next_slice(emit_len); + Ok(emit_len) +} + +#[inline(always)] +fn find_slice_mul(text: &str, markers: &[&str]) -> Option { + let range = match markers { + // Use the fast specialized `winnow::stream::FindSlice` impl for 1-3 markers. + [first] => text.find_slice(*first), + [first, second] => text.find_slice((*first, *second)), + [first, second, third] => text.find_slice((*first, *second, *third)), + // Fall back to a linear scan for 4+ markers. + _ => return markers.iter().filter_map(|marker| text.find(marker)).min(), + }; + range.map(|range| range.start) +} + /// Streaming scan state for a buffered marker search [`take_until_marker`], /// so that we don't have to rescan the whole buffered prefix when resuming. #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct MarkerScanState { +pub struct MarkerScanState { scan_start: usize, } impl MarkerScanState { - pub(super) fn reset(&mut self) { + pub fn reset(&mut self) { self.scan_start = 0; } } @@ -92,7 +131,7 @@ impl MarkerScanState { /// chunks while waiting for a closing marker. Plain `take_until` is still a /// better fit for one-shot parsers over a complete body, and for `1..` cases /// where an empty slice before the marker should be rejected. -pub(super) fn take_until_marker<'i, 'a>( +pub fn take_until_marker<'i, 'a>( marker: &'a str, state: &'a mut MarkerScanState, ) -> impl Parser, &'i str, ErrMode> + 'a { @@ -137,7 +176,7 @@ fn floor_char_boundary(text: &str, index: usize) -> usize { /// Streaming lexical state for a top-level JSON object. #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct JsonObjectScanState { +pub struct JsonObjectScanState { object_depth: usize, array_depth: usize, in_string: bool, @@ -155,7 +194,7 @@ enum JsonObjectScanPhase { impl JsonObjectScanState { /// Returns whether the top-level JSON object has closed. - pub(super) const fn complete(&self) -> bool { + pub const fn complete(&self) -> bool { matches!(self.phase, JsonObjectScanPhase::Complete) } } @@ -165,7 +204,7 @@ impl JsonObjectScanState { /// The returned length is safe to emit as raw argument text. This scans only /// lexical boundaries from `{` through the matching `}`, preserving /// malformed-but-balanced JSON without deserializing or normalizing it. -pub(super) fn take_json_object( +pub fn take_json_object( input: &mut Partial<&str>, state: &mut JsonObjectScanState, ) -> ModalResult { @@ -252,7 +291,7 @@ pub(super) fn take_json_object( } /// Parse a JSON string literal. -pub(super) fn json_str(input: &mut Partial<&str>) -> ModalResult { +pub fn json_str(input: &mut Partial<&str>) -> ModalResult { let text = **input; if text.is_empty() { return incomplete(); @@ -311,7 +350,7 @@ fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode( +pub fn parse_buffered_event( buffer: &str, parse: impl FnOnce(&mut Partial<&str>) -> ModalResult, ) -> Result> { @@ -322,7 +361,9 @@ pub(super) fn parse_buffered_event( Err(ErrMode::Incomplete(_)) => return Ok(None), Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => { // TODO: enrich context for error reporting - return Err(parsing_failed!("{}", e)); + return Err(ToolParserError::ParsingFailed { + message: e.to_string(), + }); } }; let consumed_len = input.offset_from(&checkpoint); @@ -334,7 +375,7 @@ pub(super) fn parse_buffered_event( } /// Returns an error indicating that we need more data to continue parsing. -pub(super) fn incomplete() -> ModalResult { +pub fn incomplete() -> ModalResult { Err(ErrMode::Incomplete(Needed::Unknown)) } @@ -348,7 +389,7 @@ mod tests { use super::{ JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len, - take_json_object, take_until_marker, + safe_text_len_mul, take_json_object, take_until_marker, }; #[test] @@ -406,6 +447,52 @@ mod tests { assert!(matches!(error, ErrMode::Incomplete(_))); } + #[test] + fn safe_text_len_mul_stops_before_earliest_marker() { + let mut input = Partial::new("hello<|tool_call>"); + let checkpoint = input.checkpoint(); + + let len = safe_text_len_mul(&mut input, &["<|tool_call>", ""]).unwrap(); + + assert_eq!(len, "hello".len()); + assert_eq!(input.offset_from(&checkpoint), "hello".len()); + assert_eq!(*input, "<|tool_call>"); + } + + #[test] + fn safe_text_len_mul_holds_back_longest_partial_marker() { + let mut input = Partial::new("hello<|tool"); + let checkpoint = input.checkpoint(); + + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap(); + + assert_eq!(len, "hello".len()); + assert_eq!(input.offset_from(&checkpoint), "hello".len()); + assert_eq!(*input, "<|tool"); + } + + #[test] + fn safe_text_len_mul_skips_false_same_prefix_candidate() { + let mut input = Partial::new("hello<|tool_call>"); + let checkpoint = input.checkpoint(); + + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap(); + + assert_eq!(len, "hello".len()); + assert_eq!(input.offset_from(&checkpoint), "hello".len()); + assert_eq!(*input, "<|tool_call>"); + } + + #[test] + fn safe_text_len_mul_reports_incomplete_for_only_partial_marker() { + let mut input = Partial::new("<|channel>thought"); + + let error = + safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + } + #[test] fn take_until_marker_stops_before_marker() { let mut state = MarkerScanState::default(); From 1d3f4cb3a4d0a500b479f990b8f2793d0a1a0b2f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 26 Jun 2026 08:12:38 +0800 Subject: [PATCH 009/138] [Rust Frontend] Extract renderer fixture test utilities (#46719) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 + rust/Cargo.toml | 1 + rust/src/chat/Cargo.toml | 1 + rust/src/chat/src/multimodal.rs | 15 +- .../chat/src/renderer/deepseek_v32/tests.rs | 152 +---------- .../chat/src/renderer/deepseek_v4/tests.rs | 174 +------------ rust/src/chat/src/renderer/mod.rs | 2 + rust/src/chat/src/renderer/selection.rs | 49 ++-- rust/src/chat/src/renderer/test_utils.rs | 239 ++++++++++++++++++ 9 files changed, 288 insertions(+), 346 deletions(-) create mode 100644 rust/src/chat/src/renderer/test_utils.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 70c325c152d..9ec9ac0e2da 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5722,6 +5722,7 @@ dependencies = [ "serde_json", "serde_with", "serial_test", + "strum", "subenum", "tempfile", "thiserror 2.0.18", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 601e009df07..27a758ab577 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -79,6 +79,7 @@ serde_with = "3.18.0" serial_test = { version = "3.2.0", features = ["file_locks"] } sha2 = "0.10.9" socket2 = "0.6.3" +strum = { version = "0.27.2", features = ["derive"] } subenum = "1.1.3" subtle = "2.6" task-local = "0.1.1" diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index f5860c18597..cb28b1e9c14 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -21,6 +21,7 @@ serde.workspace = true serde-json-fmt.workspace = true serde_json.workspace = true serde_with.workspace = true +strum.workspace = true subenum.workspace = true thiserror.workspace = true thiserror-ext.workspace = true diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 1b4ccc75819..024e4b63ea3 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -245,16 +245,15 @@ pub(crate) async fn finalize_rendered_prompt( return Ok((rendered.prompt, None)); } let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?; - let Prompt::Text(prompt) = rendered.prompt else { - bail_multimodal!("multimodal chat renderer must return a text prompt before expansion"); + let mut prompt_token_ids = match rendered.prompt { + Prompt::Text(prompt) => info + .context + .tokenizer() + .encode(&prompt, request.add_special_tokens) + .map_err(|error| multimodal!("{error}"))?, + Prompt::TokenIds(token_ids) => token_ids, }; let media_parts = extract_media_parts(request)?; - - let mut prompt_token_ids = info - .context - .tokenizer() - .encode(&prompt, request.add_special_tokens) - .map_err(|error| multimodal!("{error}"))?; let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?; Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 3dc3aa95795..38eddef92a2 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -1,82 +1,18 @@ -use std::fs; use std::path::PathBuf; use expect_test::{ExpectFile, expect, expect_file}; -use serde::Deserialize; use serde_json::{Value, json}; use thiserror_ext::AsReport; use super::DeepSeekV32ChatRenderer; use crate::error::Error; use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, }; use crate::{ChatRenderer, ChatRole}; -#[derive(Debug, Deserialize)] -struct FixtureRequest { - #[serde(default)] - tools: Vec, - messages: Vec, -} - -#[derive(Debug, Deserialize)] -struct FixtureTool { - function: FixtureToolFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolFunction { - name: String, - description: Option, - parameters: Value, - #[serde(default)] - strict: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "role", rename_all = "snake_case")] -enum FixtureMessage { - System { - content: String, - }, - Developer { - content: String, - #[serde(default)] - tools: Vec, - }, - User { - content: String, - }, - Assistant { - #[serde(default)] - content: String, - #[serde(default)] - reasoning_content: String, - #[serde(default)] - tool_calls: Vec, - }, - Tool { - content: String, - #[serde(default)] - tool_call_id: Option, - }, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCall { - #[serde(default)] - id: Option, - function: FixtureToolCallFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCallFunction { - name: String, - arguments: String, -} - fn render_request(request: &ChatRequest) -> String { DeepSeekV32ChatRenderer::new() .render(request) @@ -115,88 +51,14 @@ fn thinking_request(messages: Vec) -> ChatRequest { } fn fixture_request(input_name: &str) -> ChatRequest { - let fixture = fs::read_to_string(fixture_path(input_name)).unwrap(); - let fixture: FixtureRequest = serde_json::from_str(&fixture).unwrap(); - let mut request = ChatRequest { - request_id: "deepseek-v32-fixture".to_string(), - messages: fixture - .messages - .into_iter() - .enumerate() - .map(|(index, message)| match message { - FixtureMessage::System { content } => ChatMessage::system(content), - FixtureMessage::Developer { content, tools } => ChatMessage::developer( - content, - (!tools.is_empty()).then(|| to_chat_tools(&tools)), - ), - FixtureMessage::User { content } => ChatMessage::user(content), - FixtureMessage::Assistant { - content, - reasoning_content, - tool_calls, - } => { - let mut blocks = Vec::new(); - if !reasoning_content.is_empty() { - blocks.push(AssistantContentBlock::Reasoning { - text: reasoning_content, - }); - } - if !content.is_empty() { - blocks.push(AssistantContentBlock::Text { text: content }); - } - blocks.extend(tool_calls.into_iter().enumerate().map( - |(tool_index, tool_call)| { - AssistantContentBlock::ToolCall(AssistantToolCall { - id: tool_call.id.unwrap_or_else(|| { - format!("fixture-tool-call-{index}-{tool_index}") - }), - name: tool_call.function.name, - arguments: tool_call.function.arguments, - }) - }, - )); - ChatMessage::assistant_blocks(blocks) - } - FixtureMessage::Tool { - content, - tool_call_id, - } => ChatMessage::tool_response( - content, - tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")), - ), - }) - .collect(), - tools: to_chat_tools(&fixture.tools), - tool_choice: if fixture.tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, - ..ChatRequest::for_test() - }; - if matches!( - request.messages.last().map(ChatMessage::role), - Some(ChatRole::Assistant) - ) { - request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; - } - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), Value::Bool(true)); - request + fixture_chat_request(&fixture_path(input_name), deepseek_fixture_options()) } -fn to_chat_tools(tools: &[FixtureTool]) -> Vec { - tools - .iter() - .map(|tool| ChatTool { - name: tool.function.name.clone(), - description: tool.function.description.clone(), - parameters: tool.function.parameters.clone(), - strict: tool.function.strict, - }) - .collect() +fn deepseek_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + enable_thinking: true, + no_generation_prompt_when_last_assistant: true, + } } fn fixture_path(name: &str) -> PathBuf { diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 78936d8e68e..058802b8e3b 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -1,95 +1,13 @@ -use std::fs; use std::path::PathBuf; use expect_test::{ExpectFile, expect, expect_file}; -use serde::Deserialize; use serde_json::Value; use super::DeepSeekV4ChatRenderer; +use crate::ChatRenderer; use crate::event::{AssistantContentBlock, AssistantToolCall}; -use crate::request::{ - ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, ReasoningEffort, -}; -use crate::{ChatRenderer, ChatRole}; - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -enum FixtureFile { - WithTools(FixtureRequest), - MessagesOnly(Vec), -} - -#[derive(Debug, Deserialize)] -struct FixtureRequest { - #[serde(default)] - tools: Vec, - messages: Vec, -} - -impl FixtureFile { - fn into_parts(self) -> (Vec, Vec) { - match self { - Self::WithTools(req) => (req.tools, req.messages), - Self::MessagesOnly(messages) => (Vec::new(), messages), - } - } -} - -#[derive(Debug, Deserialize)] -struct FixtureTool { - function: FixtureToolFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolFunction { - name: String, - description: Option, - parameters: Value, - #[serde(default)] - strict: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "role", rename_all = "snake_case")] -enum FixtureMessage { - System { - content: String, - }, - Developer { - content: String, - #[serde(default)] - tools: Vec, - }, - User { - content: String, - }, - Assistant { - #[serde(default)] - content: String, - #[serde(default)] - reasoning_content: String, - #[serde(default)] - tool_calls: Vec, - }, - Tool { - content: String, - #[serde(default)] - tool_call_id: Option, - }, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCall { - #[serde(default)] - id: Option, - function: FixtureToolCallFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCallFunction { - name: String, - arguments: String, -} +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort}; fn render_request(request: &ChatRequest) -> String { DeepSeekV4ChatRenderer::new() @@ -101,88 +19,14 @@ fn render_request(request: &ChatRequest) -> String { } fn fixture_request(input_name: &str) -> ChatRequest { - let fixture = fs::read_to_string(fixture_path(input_name)).unwrap(); - let fixture: FixtureFile = serde_json::from_str(&fixture).unwrap(); - let (fixture_tools, fixture_messages) = fixture.into_parts(); - let mut request = ChatRequest { - request_id: "deepseek-v4-fixture".to_string(), - messages: fixture_messages - .into_iter() - .enumerate() - .map(|(index, message)| match message { - FixtureMessage::System { content } => ChatMessage::system(content), - FixtureMessage::Developer { content, tools } => ChatMessage::developer( - content, - (!tools.is_empty()).then(|| to_chat_tools(&tools)), - ), - FixtureMessage::User { content } => ChatMessage::user(content), - FixtureMessage::Assistant { - content, - reasoning_content, - tool_calls, - } => { - let mut blocks = Vec::new(); - if !reasoning_content.is_empty() { - blocks.push(AssistantContentBlock::Reasoning { - text: reasoning_content, - }); - } - if !content.is_empty() { - blocks.push(AssistantContentBlock::Text { text: content }); - } - blocks.extend(tool_calls.into_iter().enumerate().map( - |(tool_index, tool_call)| { - AssistantContentBlock::ToolCall(AssistantToolCall { - id: tool_call.id.unwrap_or_else(|| { - format!("fixture-tool-call-{index}-{tool_index}") - }), - name: tool_call.function.name, - arguments: tool_call.function.arguments, - }) - }, - )); - ChatMessage::assistant_blocks(blocks) - } - FixtureMessage::Tool { - content, - tool_call_id, - } => ChatMessage::tool_response( - content, - tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")), - ), - }) - .collect(), - tools: to_chat_tools(&fixture_tools), - tool_choice: if fixture_tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, - ..ChatRequest::for_test() - }; - if matches!( - request.messages.last().map(ChatMessage::role), - Some(ChatRole::Assistant) - ) { - request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; - } - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), Value::Bool(true)); - request + fixture_chat_request(&fixture_path(input_name), deepseek_fixture_options()) } -fn to_chat_tools(tools: &[FixtureTool]) -> Vec { - tools - .iter() - .map(|tool| ChatTool { - name: tool.function.name.clone(), - description: tool.function.description.clone(), - parameters: tool.function.parameters.clone(), - strict: tool.function.strict, - }) - .collect() +fn deepseek_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + enable_thinking: true, + no_generation_prompt_when_last_assistant: true, + } } fn fixture_path(name: &str) -> PathBuf { diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index 29e64821cf7..c4ee787c868 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -11,6 +11,8 @@ pub mod deepseek_v32; pub mod deepseek_v4; pub mod hf; mod selection; +#[cfg(test)] +mod test_utils; pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index cb22f95de0d..09bdd6b9721 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -1,10 +1,14 @@ use std::fmt; use std::str::FromStr; +use itertools::Itertools; use serde_with::{DeserializeFromStr, SerializeDisplay}; +use strum::{EnumIter, IntoEnumIterator}; /// Specify which chat renderer implementation to use. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)] +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay, EnumIter, +)] pub enum RendererSelection { /// Use model-based auto-detection. #[default] @@ -51,7 +55,8 @@ impl FromStr for RendererSelection { Ok(Self::DeepSeekV4) } else { Err(format!( - "unknown renderer `{value}` (expected one of: auto, hf, deepseek_v32, deepseek_v4)" + "unknown renderer `{value}` (expected one of: {})", + Self::iter().join(", ") )) } } @@ -70,40 +75,28 @@ impl fmt::Display for RendererSelection { #[cfg(test)] mod tests { + use std::str::FromStr as _; + + use strum::IntoEnumIterator; + use super::RendererSelection; - #[test] - fn renderer_selection_parses_known_values() { - assert_eq!( - "auto".parse::().unwrap(), - RendererSelection::Auto - ); - assert_eq!( - "hf".parse::().unwrap(), - RendererSelection::Hf - ); - assert_eq!( - "deepseek_v32".parse::().unwrap(), - RendererSelection::DeepSeekV32 - ); - assert_eq!( - "deepseek_v4".parse::().unwrap(), - RendererSelection::DeepSeekV4 - ); - } - #[test] fn renderer_selection_display_round_trips() { - for selection in [ - RendererSelection::Auto, - RendererSelection::Hf, - RendererSelection::DeepSeekV32, - RendererSelection::DeepSeekV4, - ] { + for selection in RendererSelection::iter() { assert_eq!( selection.to_string().parse::().unwrap(), selection ); } } + + #[test] + fn renderer_selection_expected_error_message() { + let err = RendererSelection::from_str("unknown").unwrap_err(); + expect_test::expect![ + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4)" + ] + .assert_eq(&err); + } } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs new file mode 100644 index 00000000000..0aab3769db4 --- /dev/null +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -0,0 +1,239 @@ +use std::fs; +use std::path::Path; + +use serde::Deserialize; +use serde_json::Value; + +use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, + GenerationPromptMode, +}; + +/// Options for constructing a [`ChatRequest`] from a fixture file. +#[derive(Debug, Clone, Copy)] +pub(crate) struct FixtureRequestOptions { + /// Whether to set the template kwarg `[enable_]thinking=true`. + pub enable_thinking: bool, + /// Whether fixtures ending in an assistant message should omit the + /// trailing generation prompt. + pub no_generation_prompt_when_last_assistant: bool, +} + +/// Read a fixture file from the given path and convert it into a [`ChatRequest`] +/// using the provided options. +pub(crate) fn fixture_chat_request(path: &Path, options: FixtureRequestOptions) -> ChatRequest { + let fixture = fs::read_to_string(path).unwrap(); + let fixture: FixtureFile = serde_json::from_str(&fixture).unwrap(); + fixture.into_request().into_chat_request(options) +} + +/// Fixture file format for chat-renderer tests. +#[derive(Debug, Deserialize)] +#[serde(untagged)] +pub(crate) enum FixtureFile { + WithRequest(FixtureRequest), + MessagesOnly(Vec), +} + +#[derive(Debug, Deserialize)] +pub(crate) struct FixtureRequest { + #[serde(default)] + tools: Vec, + messages: Vec, + add_generation_prompt: Option, +} + +impl FixtureFile { + fn into_request(self) -> FixtureRequest { + match self { + Self::WithRequest(request) => request, + Self::MessagesOnly(messages) => FixtureRequest { + tools: Vec::new(), + messages, + add_generation_prompt: None, + }, + } + } +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "role", rename_all = "snake_case")] +pub(crate) enum FixtureMessage { + System { + content: FixtureContent, + }, + Developer { + content: FixtureContent, + #[serde(default)] + tools: Vec, + }, + User { + content: FixtureContent, + }, + Assistant { + #[serde(default)] + content: String, + #[serde(default)] + reasoning_content: String, + #[serde(default)] + tool_calls: Vec, + }, + Tool { + content: FixtureContent, + #[serde(default)] + tool_call_id: Option, + }, +} + +#[derive(Debug, Deserialize)] +#[serde(untagged)] +pub(crate) enum FixtureContent { + Text(String), + Parts(Vec), +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub(crate) enum FixtureContentPart { + Text { text: String }, + ImageUrl { image_url: String }, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct FixtureTool { + function: FixtureToolFunction, +} + +#[derive(Debug, Deserialize)] +struct FixtureToolFunction { + name: String, + description: Option, + parameters: Value, + #[serde(default)] + strict: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct FixtureToolCall { + #[serde(default)] + id: Option, + function: FixtureToolCallFunction, +} + +#[derive(Debug, Deserialize)] +struct FixtureToolCallFunction { + name: String, + arguments: String, +} + +impl FixtureRequest { + fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest { + let mut request = ChatRequest { + request_id: "renderer-fixture".to_string(), + messages: self + .messages + .into_iter() + .enumerate() + .map(|(index, message)| fixture_message_to_chat_message(index, message)) + .collect(), + tools: to_chat_tools(&self.tools), + tool_choice: if self.tools.is_empty() { + ChatToolChoice::None + } else { + ChatToolChoice::Auto + }, + ..ChatRequest::for_test() + }; + + if options.no_generation_prompt_when_last_assistant + && matches!(request.messages.last(), Some(ChatMessage::Assistant { .. })) + { + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + } + if self.add_generation_prompt == Some(false) { + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + } + if options.enable_thinking { + for key in ["thinking", "enable_thinking"] { + request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true)); + } + } + + request + } +} + +fn fixture_message_to_chat_message(index: usize, message: FixtureMessage) -> ChatMessage { + match message { + FixtureMessage::System { content } => ChatMessage::system(to_chat_content(content)), + FixtureMessage::Developer { content, tools } => ChatMessage::developer( + to_chat_content(content), + (!tools.is_empty()).then(|| to_chat_tools(&tools)), + ), + FixtureMessage::User { content } => ChatMessage::user(to_chat_content(content)), + FixtureMessage::Assistant { + content, + reasoning_content, + tool_calls, + } => { + let mut blocks = Vec::new(); + if !reasoning_content.is_empty() { + blocks.push(AssistantContentBlock::Reasoning { + text: reasoning_content, + }); + } + if !content.is_empty() { + blocks.push(AssistantContentBlock::Text { text: content }); + } + blocks.extend( + tool_calls.into_iter().enumerate().map(|(tool_index, tool_call)| { + AssistantContentBlock::ToolCall(AssistantToolCall { + id: tool_call + .id + .unwrap_or_else(|| format!("fixture-tool-call-{index}-{tool_index}")), + name: tool_call.function.name, + arguments: tool_call.function.arguments, + }) + }), + ); + ChatMessage::assistant_blocks(blocks) + } + FixtureMessage::Tool { + content, + tool_call_id, + } => ChatMessage::tool_response( + to_chat_content(content), + tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")), + ), + } +} + +fn to_chat_content(content: FixtureContent) -> ChatContent { + match content { + FixtureContent::Text(text) => ChatContent::Text(text), + FixtureContent::Parts(parts) => ChatContent::Parts( + parts + .into_iter() + .map(|part| match part { + FixtureContentPart::Text { text } => ChatContentPart::text(text), + FixtureContentPart::ImageUrl { image_url } => { + ChatContentPart::image_url(image_url) + } + }) + .collect(), + ), + } +} + +fn to_chat_tools(tools: &[FixtureTool]) -> Vec { + tools + .iter() + .map(|tool| ChatTool { + name: tool.function.name.clone(), + description: tool.function.description.clone(), + parameters: tool.function.parameters.clone(), + strict: tool.function.strict, + }) + .collect() +} From ae7c8ec223e4d6bdfdaed6c8bb58e54b44d4ccaf Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 26 Jun 2026 08:19:44 +0800 Subject: [PATCH 010/138] [Rust Frontend] Switch `rustls` to `native-tls`/OpenSSL (#46696) Signed-off-by: Bugen Zhao --- .../scripts/run-rust-frontend-cargo-ci.sh | 18 + rust/Cargo.lock | 438 +----------------- rust/Cargo.toml | 12 +- rust/deny.toml | 15 + rust/src/text/Cargo.toml | 1 + rust/src/tokenizer/Cargo.toml | 2 + rust/src/tokenizer/benches/hf.rs | 20 +- rust/src/tokenizer/benches/tiktoken.rs | 24 +- 8 files changed, 89 insertions(+), 441 deletions(-) create mode 100644 rust/deny.toml diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index 4b4272762a1..42ab1fb543b 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -90,6 +90,16 @@ install_cargo_sort() { cargo binstall --no-confirm cargo-sort } +install_cargo_deny() { + if command -v cargo-deny >/dev/null 2>&1; then + return + fi + + log_section "Installing cargo-deny" + install_cargo_binstall + cargo binstall --no-confirm cargo-deny +} + install_cargo_nextest() { if command -v cargo-nextest >/dev/null 2>&1; then return @@ -142,6 +152,7 @@ PY run_style_clippy() { install_cargo_sort + install_cargo_deny log_section "Checking Rust formatting" cargo fmt --manifest-path rust/Cargo.toml --all -- --check @@ -149,6 +160,13 @@ run_style_clippy() { log_section "Checking Cargo.toml ordering" cargo sort --workspace --check rust + log_section "Checking Rust dependency bans" + cargo deny \ + --manifest-path rust/Cargo.toml \ + check \ + --config rust/deny.toml \ + bans + log_section "Running clippy" cargo clippy \ --manifest-path rust/Cargo.toml \ diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 9ec9ac0e2da..52635c75e8f 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -454,12 +454,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "base64ct" -version = "1.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" - [[package]] name = "bit-set" version = "0.5.3" @@ -631,12 +625,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" -[[package]] -name = "cfg_aliases" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" - [[package]] name = "chrono" version = "0.4.44" @@ -754,19 +742,6 @@ dependencies = [ "crossbeam-utils", ] -[[package]] -name = "console" -version = "0.15.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" -dependencies = [ - "encode_unicode", - "libc", - "once_cell", - "unicode-width", - "windows-sys 0.59.0", -] - [[package]] name = "console" version = "0.16.2" @@ -786,35 +761,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" -[[package]] -name = "cookie" -version = "0.18.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" -dependencies = [ - "percent-encoding", - "time", - "version_check", -] - -[[package]] -name = "cookie_store" -version = "0.22.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" -dependencies = [ - "cookie", - "document-features", - "idna", - "indexmap 2.13.0", - "log", - "serde", - "serde_derive", - "serde_json", - "time", - "url", -] - [[package]] name = "core-foundation" version = "0.9.4" @@ -1038,16 +984,6 @@ dependencies = [ "serde", ] -[[package]] -name = "der" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71fd89660b2dc699704064e59e9dba0147b903e85319429e131620d022be411b" -dependencies = [ - "pem-rfc7468", - "zeroize", -] - [[package]] name = "deranged" version = "0.5.8" @@ -1393,13 +1329,12 @@ dependencies = [ [[package]] name = "fastokens" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "796a262ed47d1458a4b40d0ed831c927e6f54d5b9c1de2683bb4ac9b04f4c7cc" +checksum = "8728655e193e0d08d7a95d63cf1fdb9b768d282cab0a112ecb006615bae9f067" dependencies = [ "daachorse", "fancy-regex 0.17.0", - "hf-hub 0.4.3", "icu_normalizer", "memchr", "pcre2", @@ -1643,10 +1578,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", - "js-sys", "libc", "wasi", - "wasm-bindgen", ] [[package]] @@ -1765,26 +1698,6 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" -[[package]] -name = "hf-hub" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" -dependencies = [ - "dirs", - "http", - "indicatif 0.17.11", - "libc", - "log", - "rand 0.9.2", - "reqwest", - "serde", - "serde_json", - "thiserror 2.0.18", - "ureq 2.12.1", - "windows-sys 0.60.2", -] - [[package]] name = "hf-hub" version = "0.5.0" @@ -1793,11 +1706,9 @@ checksum = "aef3982638978efa195ff11b305f51f1f22f4f0a6cabee7af79b383ebee6a213" dependencies = [ "dirs", "futures", - "http", - "indicatif 0.18.4", + "indicatif", "libc", "log", - "native-tls", "num_cpus", "rand 0.9.2", "reqwest", @@ -1805,7 +1716,6 @@ dependencies = [ "serde_json", "thiserror 2.0.18", "tokio", - "ureq 3.3.0", "windows-sys 0.61.2", ] @@ -1902,12 +1812,10 @@ dependencies = [ "hyper", "hyper-util", "rustls", - "rustls-native-certs", "rustls-pki-types", "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.6", ] [[package]] @@ -2168,26 +2076,13 @@ dependencies = [ "serde_core", ] -[[package]] -name = "indicatif" -version = "0.17.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" -dependencies = [ - "console 0.15.11", - "number_prefix", - "portable-atomic", - "unicode-width", - "web-time", -] - [[package]] name = "indicatif" version = "0.18.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" dependencies = [ - "console 0.16.2", + "console", "portable-atomic", "unicode-width", "unit-prefix", @@ -2421,7 +2316,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.5.0" -source = "git+https://github.com/vllm-project/llm-multimodal?rev=5b558989844d1c7af3e43d0f604069ffd9c06320#5b558989844d1c7af3e43d0f604069ffd9c06320" +source = "git+https://github.com/vllm-project/llm-multimodal?rev=046b669bd1c4faa2a7e05344d8cbf7b2befb37d5#046b669bd1c4faa2a7e05344d8cbf7b2befb37d5" dependencies = [ "base64 0.22.1", "blake3", @@ -2462,12 +2357,6 @@ dependencies = [ "imgref", ] -[[package]] -name = "lru-slab" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" - [[package]] name = "macro_rules_attribute" version = "0.2.2" @@ -2881,12 +2770,6 @@ dependencies = [ "libc", ] -[[package]] -name = "number_prefix" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" - [[package]] name = "once_cell" version = "1.21.3" @@ -2930,8 +2813,7 @@ checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" [[package]] name = "openai-harmony" version = "0.0.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e77e82af451fc95deeb728a40b84db8ee82d341e136c268de415123a560b9b72" +source = "git+https://github.com/Inferact/openai-harmony?rev=cfbadbc66f3158692bfeefa961e363aa7a6b9708#cfbadbc66f3158692bfeefa961e363aa7a6b9708" dependencies = [ "anyhow", "base64 0.22.1", @@ -3093,15 +2975,6 @@ dependencies = [ "pkg-config", ] -[[package]] -name = "pem-rfc7468" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9" -dependencies = [ - "base64ct", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -3553,61 +3426,6 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" -[[package]] -name = "quinn" -version = "0.11.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" -dependencies = [ - "bytes", - "cfg_aliases", - "pin-project-lite", - "quinn-proto", - "quinn-udp", - "rustc-hash 2.1.1", - "rustls", - "socket2", - "thiserror 2.0.18", - "tokio", - "tracing", - "web-time", -] - -[[package]] -name = "quinn-proto" -version = "0.11.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" -dependencies = [ - "bytes", - "getrandom 0.3.4", - "lru-slab", - "rand 0.9.2", - "ring", - "rustc-hash 2.1.1", - "rustls", - "rustls-pki-types", - "slab", - "thiserror 2.0.18", - "tinyvec", - "tracing", - "web-time", -] - -[[package]] -name = "quinn-udp" -version = "0.5.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" -dependencies = [ - "cfg_aliases", - "libc", - "once_cell", - "socket2", - "tracing", - "windows-sys 0.60.2", -] - [[package]] name = "quote" version = "1.0.45" @@ -3880,9 +3698,6 @@ dependencies = [ "native-tls", "percent-encoding", "pin-project-lite", - "quinn", - "rustls", - "rustls-native-certs", "rustls-pki-types", "serde", "serde_json", @@ -3890,7 +3705,6 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-native-tls", - "tokio-rustls", "tokio-util", "tower", "tower-http", @@ -3900,7 +3714,6 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots 1.0.6", ] [[package]] @@ -4039,34 +3852,19 @@ version = "0.23.37" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" dependencies = [ - "log", "once_cell", - "ring", "rustls-pki-types", "rustls-webpki", "subtle", "zeroize", ] -[[package]] -name = "rustls-native-certs" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" -dependencies = [ - "openssl-probe", - "rustls-pki-types", - "schannel", - "security-framework", -] - [[package]] name = "rustls-pki-types" version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" dependencies = [ - "web-time", "zeroize", ] @@ -4600,17 +4398,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "socks" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" -dependencies = [ - "byteorder", - "libc", - "winapi", -] - [[package]] name = "spm_precompiled" version = "0.1.4" @@ -4976,21 +4763,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "tinyvec" -version = "1.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" -dependencies = [ - "tinyvec_macros", -] - -[[package]] -name = "tinyvec_macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" - [[package]] name = "tokenizers" version = "0.22.2" @@ -5004,7 +4776,7 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", - "indicatif 0.18.4", + "indicatif", "itertools 0.14.0", "log", "macro_rules_attribute", @@ -5529,61 +5301,6 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" -[[package]] -name = "ureq" -version = "2.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" -dependencies = [ - "base64 0.22.1", - "flate2", - "log", - "once_cell", - "rustls", - "rustls-pki-types", - "serde", - "serde_json", - "socks", - "url", - "webpki-roots 0.26.11", -] - -[[package]] -name = "ureq" -version = "3.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" -dependencies = [ - "base64 0.22.1", - "cookie_store", - "der", - "flate2", - "log", - "native-tls", - "percent-encoding", - "rustls", - "rustls-pki-types", - "serde", - "serde_json", - "socks", - "ureq-proto", - "utf8-zero", - "webpki-root-certs", - "webpki-roots 1.0.6", -] - -[[package]] -name = "ureq-proto" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" -dependencies = [ - "base64 0.22.1", - "http", - "httparse", - "log", -] - [[package]] name = "url" version = "2.5.8" @@ -5608,12 +5325,6 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" -[[package]] -name = "utf8-zero" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" - [[package]] name = "utf8_iter" version = "1.0.4" @@ -5951,8 +5662,9 @@ dependencies = [ "enum-as-inner", "expect-test", "futures", - "hf-hub 0.5.0", + "hf-hub", "itertools 0.14.0", + "reqwest", "serde", "serde_json", "serde_with", @@ -5975,7 +5687,8 @@ dependencies = [ "base64 0.22.1", "criterion", "fastokens", - "hf-hub 0.5.0", + "hf-hub", + "reqwest", "riptoken", "rustc-hash 1.1.0", "serde", @@ -5986,6 +5699,7 @@ dependencies = [ "thiserror-ext", "tiktoken-rs 0.9.1", "tokenizers", + "tokio", "tracing", ] @@ -6169,33 +5883,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "webpki-root-certs" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804f18a4ac2676ffb4e8b5b5fa9ae38af06df08162314f96a68d2a363e21a8ca" -dependencies = [ - "rustls-pki-types", -] - -[[package]] -name = "webpki-roots" -version = "0.26.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" -dependencies = [ - "webpki-roots 1.0.6", -] - -[[package]] -name = "webpki-roots" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed" -dependencies = [ - "rustls-pki-types", -] - [[package]] name = "weezl" version = "0.1.12" @@ -6320,25 +6007,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -6356,31 +6025,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -6389,96 +6041,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - [[package]] name = "winnow" version = "1.0.2" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 27a758ab577..4f6322e7ada 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -23,7 +23,7 @@ license = "Apache-2.0" [workspace.dependencies] anyhow = "1.0.100" arc-swap = "1.9.0" -async-openai = "0.33.1" +async-openai = { version = "0.33.1", default-features = false, features = ["native-tls"] } async-trait = "0.1.89" asynk-strim-attr = "0.1.0" axum = "0.8.8" @@ -37,22 +37,22 @@ easy-ext = "1.0.3" educe = "0.6.0" enum-as-inner = "0.7.0" expect-test = "1.5.1" -fastokens = "0.2.0" +fastokens = { version = "0.2.1", default-features = false } futures = "0.3.31" half = { version = "2.7.1", features = ["bytemuck"] } hex = "0.4.3" -hf-hub = { version = "0.5.0", features = ["tokio"] } +hf-hub = { version = "0.5.0", default-features = false, features = ["tokio"] } http-body = "1.0.1" indexmap = "2.13.0" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" } +llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "046b669bd1c4faa2a7e05344d8cbf7b2befb37d5" } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } ndarray = { version = "0.16.1", features = ["serde"] } -openai-harmony = "0.0.8" +openai-harmony = { git = "https://github.com/Inferact/openai-harmony", rev = "cfbadbc66f3158692bfeefa961e363aa7a6b9708", default-features = false, features = ["native-tls"] } openai-protocol = "1.6.0" parking_lot = "0.12.5" paste = "1.0.15" @@ -64,7 +64,7 @@ pyo3 = "0.28.3" pythonize = "0.28.0" rand = "0.9.2" reasoning-parser = "1.2.2" -reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] } +reqwest = { version = "0.12.8", default-features = false, features = ["native-tls"] } riptoken = { version = "0.3.0", default-features = false } rmp-serde = "1.3.1" rmpv = { version = "1.3.1", features = ["with-serde"] } diff --git a/rust/deny.toml b/rust/deny.toml new file mode 100644 index 00000000000..25bd8e3831a --- /dev/null +++ b/rust/deny.toml @@ -0,0 +1,15 @@ +[bans] +multiple-versions = "allow" + +deny = [ + # TLS / crypto provider + # We prefer the system's TLS (e.g. OpenSSL) over Rust implementations. + { name = "rustls" }, + { name = "ring" }, + { name = "aws-lc-rs" }, + { name = "aws-lc-sys" }, + { name = "s2n-tls" }, + { name = "s2n-tls-sys" }, + { name = "boring" }, + { name = "boring-sys" }, +] diff --git a/rust/src/text/Cargo.toml b/rust/src/text/Cargo.toml index 8be9ee78764..7ed02c07fca 100644 --- a/rust/src/text/Cargo.toml +++ b/rust/src/text/Cargo.toml @@ -12,6 +12,7 @@ enum-as-inner.workspace = true futures.workspace = true hf-hub.workspace = true itertools.workspace = true +reqwest.workspace = true serde.workspace = true serde_json.workspace = true serde_with.workspace = true diff --git a/rust/src/tokenizer/Cargo.toml b/rust/src/tokenizer/Cargo.toml index 786c46f4031..7b54676f66b 100644 --- a/rust/src/tokenizer/Cargo.toml +++ b/rust/src/tokenizer/Cargo.toml @@ -21,7 +21,9 @@ tracing.workspace = true [dev-dependencies] criterion.workspace = true hf-hub.workspace = true +reqwest.workspace = true tempfile.workspace = true +tokio.workspace = true [[bench]] name = "hf" diff --git a/rust/src/tokenizer/benches/hf.rs b/rust/src/tokenizer/benches/hf.rs index 9bf37778089..950c4784a74 100644 --- a/rust/src/tokenizer/benches/hf.rs +++ b/rust/src/tokenizer/benches/hf.rs @@ -1,5 +1,6 @@ use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main}; -use hf_hub::api::sync::ApiBuilder; +use hf_hub::api::tokio::ApiBuilder; +use tokio::runtime::Runtime; use vllm_tokenizer::{HuggingFaceTokenizer, Tokenizer}; const MODEL_ID: &str = "Qwen/Qwen3.5-0.8B"; @@ -55,13 +56,16 @@ impl BenchFixture { } fn tokenizer_json() -> std::path::PathBuf { - ApiBuilder::from_env() - .with_progress(false) - .build() - .expect("build hf-hub api") - .model(MODEL_ID.to_string()) - .get("tokenizer.json") - .expect("fetch tokenizer.json from hf-hub") + Runtime::new().expect("build tokio runtime").block_on(async { + ApiBuilder::from_env() + .with_progress(false) + .build() + .expect("build hf-hub api") + .model(MODEL_ID.to_string()) + .get("tokenizer.json") + .await + .expect("fetch tokenizer.json from hf-hub") + }) } fn bench_encode(c: &mut Criterion) { diff --git a/rust/src/tokenizer/benches/tiktoken.rs b/rust/src/tokenizer/benches/tiktoken.rs index 54b9805f01a..6540adf486d 100644 --- a/rust/src/tokenizer/benches/tiktoken.rs +++ b/rust/src/tokenizer/benches/tiktoken.rs @@ -1,5 +1,6 @@ use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main}; -use hf_hub::api::sync::ApiBuilder; +use hf_hub::api::tokio::ApiBuilder; +use tokio::runtime::Runtime; use vllm_tokenizer::{TiktokenTokenizer, Tokenizer}; const MODEL_ID: &str = "moonshotai/Kimi-K2.5"; @@ -52,15 +53,18 @@ impl BenchFixture { } fn tiktoken_model() -> std::path::PathBuf { - let repo = ApiBuilder::from_env() - .with_progress(false) - .build() - .expect("build hf-hub api") - .model(MODEL_ID.to_string()); - repo.get("config.json").expect("fetch config.json from hf-hub"); - repo.get("tokenizer_config.json") - .expect("fetch tokenizer_config.json from hf-hub"); - repo.get("tiktoken.model").expect("fetch tiktoken.model from hf-hub") + Runtime::new().expect("build tokio runtime").block_on(async { + let repo = ApiBuilder::from_env() + .with_progress(false) + .build() + .expect("build hf-hub api") + .model(MODEL_ID.to_string()); + repo.get("config.json").await.expect("fetch config.json from hf-hub"); + repo.get("tokenizer_config.json") + .await + .expect("fetch tokenizer_config.json from hf-hub"); + repo.get("tiktoken.model").await.expect("fetch tiktoken.model from hf-hub") + }) } fn bench_encode(c: &mut Criterion) { From ad28d605e6db88b7236977517799b5088076f209 Mon Sep 17 00:00:00 2001 From: Mike G Date: Thu, 25 Jun 2026 17:46:28 -0700 Subject: [PATCH 011/138] [Bugfix] Default tie_weights to sharing the weight (fix tied quantized embeddings, e.g. ModelOpt Gemma4) (#45544) Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com> Co-authored-by: Michael Goin --- .../layers/quantization/base_config.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 7bc5d16be73..9b18bdc132e 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -48,11 +48,18 @@ class QuantizeMethodBase(ABC): raise NotImplementedError # Not required functions - def tie_weights(self, layer: torch.nn.Module, *args, **kwargs): - """Tie layer's weights for the layer from another layer/tensors. + def tie_weights(self, layer: torch.nn.Module, embed_tokens: torch.nn.Module): + """Tie ``layer``'s weight to ``embed_tokens``' weight. + + The default shares the weight tensor, which is the standard behavior for + tied word embeddings and matches what ``ParallelLMHead.tie_weights`` did + directly before quantization methods became responsible for it. + Quantization methods that need special weight handling (e.g. repacked + weights) override this. Expects create_weights to have been called before on the layer.""" - raise NotImplementedError + layer.weight = embed_tokens.weight + return layer def process_weights_after_loading(self, layer: nn.Module) -> None: """Process the weight after loading. From 32bb3195f0b93f6971781479591f7a6ee666e7dc Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 25 Jun 2026 18:04:06 -0700 Subject: [PATCH 012/138] [ModelRunner V2] Bound memory for large logprobs requests (#46746) Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/sample/logprob.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/vllm/v1/worker/gpu/sample/logprob.py b/vllm/v1/worker/gpu/sample/logprob.py index 0028e8c3a9d..cb2cf1a590e 100644 --- a/vllm/v1/worker/gpu/sample/logprob.py +++ b/vllm/v1/worker/gpu/sample/logprob.py @@ -9,6 +9,9 @@ from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsTensors from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor +# Upper bound on the topk kernel's per-iteration gather width. +_MAX_TOPK_BLOCK = 1024 + @triton.jit def _topk_log_softmax_kernel( @@ -19,7 +22,7 @@ def _topk_log_softmax_kernel( topk, vocab_size, BLOCK_SIZE: tl.constexpr, - PADDED_TOPK: tl.constexpr, + TOPK_BLOCK_SIZE: tl.constexpr, ): req_idx = tl.program_id(0).to(tl.int64) row_ptr = logits_ptr + req_idx * logits_stride @@ -42,14 +45,16 @@ def _topk_log_softmax_kernel( se += tl.sum(e) lse = tl.log(se) - k_offset = tl.arange(0, PADDED_TOPK) - k_mask = k_offset < topk - topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0) - - logits = tl.load(row_ptr + topk_ids, mask=k_mask) - logits = logits.to(tl.float32) - o = logits - max_val - lse - tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask) + for j in range(0, topk, TOPK_BLOCK_SIZE): + k_offset = j + tl.arange(0, TOPK_BLOCK_SIZE) + k_mask = k_offset < topk + topk_ids = tl.load( + topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0 + ) + logits = tl.load(row_ptr + topk_ids, mask=k_mask) + logits = logits.to(tl.float32) + o = logits - max_val - lse + tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask) @triton.jit @@ -85,6 +90,9 @@ def compute_token_logprobs( token_ids = token_ids.to(torch.int64) num_logprobs = token_ids.shape[1] logprobs = logits.new_empty((batch_size, num_logprobs), dtype=torch.float32) + # Cap the kernel's per-iteration width so very large num_logprobs requests + # stream the gather in bounded-size chunks, avoiding excessive mem use. + topk_block_size = min(triton.next_power_of_2(num_logprobs), _MAX_TOPK_BLOCK) _topk_log_softmax_kernel[(batch_size,)]( logprobs, logits, @@ -93,7 +101,7 @@ def compute_token_logprobs( num_logprobs, vocab_size, BLOCK_SIZE=1024, # type: ignore - PADDED_TOPK=triton.next_power_of_2(num_logprobs), + TOPK_BLOCK_SIZE=topk_block_size, ) return logprobs From cc7981599eac6d6ed6d08c07f9ec47d771969712 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:09:56 -0400 Subject: [PATCH 013/138] [Refactor] Remove dead kernel code (#46405) Signed-off-by: yewentao256 --- csrc/custom_all_reduce_test.cu | 361 --------------------------------- csrc/ops.h | 6 - vllm/_custom_ops.py | 64 ------ 3 files changed, 431 deletions(-) delete mode 100644 csrc/custom_all_reduce_test.cu diff --git a/csrc/custom_all_reduce_test.cu b/csrc/custom_all_reduce_test.cu deleted file mode 100644 index f7f0823465d..00000000000 --- a/csrc/custom_all_reduce_test.cu +++ /dev/null @@ -1,361 +0,0 @@ -/** - * This is a standalone test for custom allreduce. - * To compile, make sure you have MPI and NCCL installed in your system. - * export MPI_HOME=XXX - * nvcc -O2 -arch=native -std=c++17 custom_all_reduce_test.cu -o - * custom_all_reduce_test -lnccl -I${MPI_HOME}/include -lmpi - * - * Warning: this C++ test is not designed to be very readable and was used - * during the rapid prototyping process. - * - * To run: - * mpirun --allow-run-as-root -np 8 ./custom_all_reduce_test - */ -#include -#include -#include -#include - -#include -#include - -#include "cuda_profiler_api.h" -#include "custom_all_reduce.cuh" -#include "mpi.h" -#ifdef USE_ROCM - #include -typedef __hip_bfloat16 nv_bfloat16; - #include "rccl/rccl.h" - #include "custom_all_reduce_hip.cuh" -#else - #include "nccl.h" - #include "custom_all_reduce.cuh" -#endif - -#define MPICHECK(cmd) \ - do { \ - int e = cmd; \ - if (e != MPI_SUCCESS) { \ - printf("Failed: MPI error %s:%d '%d'\n", __FILE__, __LINE__, e); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -#define NCCLCHECK(cmd) \ - do { \ - ncclResult_t r = cmd; \ - if (r != ncclSuccess) { \ - printf("Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, \ - ncclGetErrorString(r)); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -#ifdef USE_ROCM -__global__ void dummy_kernel() { - for (int i = 0; i < 100; i++) { - uint64_t start = wall_clock64(); - uint64_t cycles_elapsed; - do { - cycles_elapsed = wall_clock64() - start; - } while (cycles_elapsed < 100); - } - for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms -} -#else -__global__ void dummy_kernel() { - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms - #else - for (int i = 0; i < 100; i++) { - long long int start = clock64(); - while (clock64() - start < 150000000); // approximately 98.4ms on P40 - } - #endif -} -#endif - -template -__global__ void set_data(T* data, int size, int myRank) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - data[idx] = myRank * 0.11f; - } -} - -template -__global__ void convert_data(const T* data1, const T* data2, double* fdata1, - double* fdata2, int size) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - fdata1[idx] = data1[idx]; - fdata2[idx] = data2[idx]; - } -} - -__global__ void init_rand(curandState_t* state, int size, int nRanks) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - for (int i = 0; i < nRanks; i++) { - curand_init(i + 1, idx, 0, &state[idx * nRanks + i]); - } - } -} - -template -__global__ void gen_data(curandState_t* state, T* data, double* ground_truth, - int myRank, int nRanks, int size) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - double sum = 0.0; - for (int i = 0; i < nRanks; i++) { - double val = curand_uniform_double(&state[idx * nRanks + i]) * 4; - T hval = val; // downcast first - sum += static_cast(hval); - if (i == myRank) data[idx] = hval; - } - ground_truth[idx] = sum; - } -} - -template -void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit, - int data_size, bool performance_test) { - T* result; - cudaStream_t stream; - CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); - CUDACHECK(cudaMalloc(&result, data_size * sizeof(T))); - CUDACHECK(cudaMemset(result, 0, data_size * sizeof(T))); - - cudaIpcMemHandle_t self_data_handle; - cudaIpcMemHandle_t data_handles[8]; - vllm::Signal* buffer; - T* self_data_copy; - /** - * Allocate IPC buffer - * - * The first section is a temporary buffer for storing intermediate allreduce - * results, if a particular algorithm requires it. The second section is for - * the input to the allreduce. The actual API takes the input pointer as an - * argument (that is, they can and usually should be allocated separately). - * But since the input pointers and the temporary buffer all require IPC - * registration, they are allocated and registered together in the test for - * convenience. - */ -#ifdef USE_ROCM - CUDACHECK(hipExtMallocWithFlags( - (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal), - hipDeviceMallocUncached)); -#else - CUDACHECK( - cudaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal))); -#endif - CUDACHECK( - cudaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm::Signal))); - CUDACHECK(cudaMalloc(&self_data_copy, data_size * sizeof(T))); - CUDACHECK(cudaIpcGetMemHandle(&self_data_handle, buffer)); - - MPICHECK(MPI_Allgather(&self_data_handle, sizeof(cudaIpcMemHandle_t), - MPI_BYTE, data_handles, sizeof(cudaIpcMemHandle_t), - MPI_BYTE, MPI_COMM_WORLD)); - - void* rank_data; - size_t rank_data_sz = 16 * 1024 * 1024; - CUDACHECK(cudaMalloc(&rank_data, rank_data_sz)); - vllm::Signal* ipc_ptrs[8]; - for (int i = 0; i < nRanks; i++) { - if (i == myRank) - ipc_ptrs[i] = buffer; - else - CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptrs[i], data_handles[i], - cudaIpcMemLazyEnablePeerAccess)); - } - vllm::CustomAllreduce fa(ipc_ptrs, rank_data, rank_data_sz, myRank, nRanks); - auto* self_data = - reinterpret_cast(reinterpret_cast(buffer) + - sizeof(vllm::Signal) + data_size * sizeof(T)); - // hack buffer registration - { - void* data[8]; - for (int i = 0; i < nRanks; i++) { - data[i] = - ((char*)ipc_ptrs[i]) + sizeof(vllm::Signal) + data_size * sizeof(T); - } - fa.register_buffer(data); - } - - double* ground_truth; - CUDACHECK(cudaMallocHost(&ground_truth, data_size * sizeof(double))); - curandState_t* states; - CUDACHECK(cudaMalloc(&states, sizeof(curandState_t) * nRanks * data_size)); - init_rand<<<108, 1024, 0, stream>>>(states, data_size, nRanks); - gen_data<<<108, 1024, 0, stream>>>(states, self_data, ground_truth, myRank, - nRanks, data_size); - CUDACHECK(cudaMemcpyAsync(self_data_copy, self_data, data_size * sizeof(T), - cudaMemcpyDeviceToDevice, stream)); - cudaEvent_t start, stop; - CUDACHECK(cudaEventCreate(&start)); - CUDACHECK(cudaEventCreate(&stop)); - - ncclDataType_t ncclDtype; - if (std::is_same::value) { - ncclDtype = ncclFloat16; - } else if (std::is_same::value) { - ncclDtype = ncclBfloat16; - } else { - ncclDtype = ncclFloat; - } - double *nccl_result, *my_result; - CUDACHECK(cudaMallocHost(&nccl_result, data_size * sizeof(double))); - CUDACHECK(cudaMallocHost(&my_result, data_size * sizeof(double))); - if (performance_test) { - dummy_kernel<<<1, 1, 0, stream>>>(); - constexpr int warmup_iters = 5; - constexpr int num_iters = 100; - // warmup - for (int i = 0; i < warmup_iters; i++) { - NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum, - comm, stream)); - } - CUDACHECK(cudaEventRecord(start, stream)); - for (int i = 0; i < num_iters; i++) { - NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum, - comm, stream)); - } - CUDACHECK(cudaEventRecord(stop, stream)); - CUDACHECK(cudaStreamSynchronize(stream)); - float allreduce_ms = 0; - cudaEventElapsedTime(&allreduce_ms, start, stop); - - dummy_kernel<<<1, 1, 0, stream>>>(); - // warm up - for (int i = 0; i < warmup_iters; i++) { - fa.allreduce(stream, self_data, result, data_size, threads, - block_limit); - } - CUDACHECK(cudaEventRecord(start, stream)); - for (int i = 0; i < num_iters; i++) { - fa.allreduce(stream, self_data, result, data_size, threads, - block_limit); - } - CUDACHECK(cudaEventRecord(stop, stream)); - CUDACHECK(cudaStreamSynchronize(stream)); - - float duration_ms = 0; - cudaEventElapsedTime(&duration_ms, start, stop); - if (myRank == 0) - printf( - "Rank %d done, nGPUs:%d, sz (kb): %d, %d, %d, my time:%.2fus, nccl " - "time:%.2fus\n", - myRank, nRanks, data_size * sizeof(T) / 1024, threads, block_limit, - duration_ms * 1e3 / num_iters, allreduce_ms * 1e3 / num_iters); - - // And wait for all the queued up work to complete - CUDACHECK(cudaStreamSynchronize(stream)); - - NCCLCHECK(ncclAllReduce(self_data_copy, self_data, data_size, ncclDtype, - ncclSum, comm, stream)); - - convert_data<<<108, 1024, 0, stream>>>(self_data, result, nccl_result, - my_result, data_size); - CUDACHECK(cudaStreamSynchronize(stream)); - - for (unsigned long j = 0; j < data_size; j++) { - auto diff = abs(nccl_result[j] - my_result[j]); - if (diff >= 4e-2) { - printf("Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n", - myRank, j, nccl_result[j], my_result[j], ground_truth[j]); - break; - } - } - long double nccl_diffs = 0.0; - long double my_diffs = 0.0; - for (int j = 0; j < data_size; j++) { - nccl_diffs += abs(nccl_result[j] - ground_truth[j]); - my_diffs += abs(my_result[j] - ground_truth[j]); - } - if (myRank == 0) - std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size - << " me: " << my_diffs / data_size << std::endl; - } else { - for (int i = 0; i < 100; i++) { - fa.allreduce(stream, self_data, result, data_size, threads, - block_limit); - CUDACHECK(cudaStreamSynchronize(stream)); - NCCLCHECK(ncclAllReduce(self_data, self_data_copy, data_size, ncclDtype, - ncclSum, comm, stream)); - convert_data<<<108, 1024, 0, stream>>>( - self_data_copy, result, nccl_result, my_result, data_size); - CUDACHECK(cudaStreamSynchronize(stream)); - - for (unsigned long j = 0; j < data_size; j++) { - auto diff = abs(nccl_result[j] - my_result[j]); - if (diff >= 4e-2) { - printf( - "Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n", - myRank, j, nccl_result[j], my_result[j], ground_truth[j]); - break; - } - } - } - if (myRank == 0) - printf("Test passed: nGPUs:%d, sz (kb): %d, %d, %d\n", nRanks, - data_size * sizeof(T) / 1024, threads, block_limit); - // long double nccl_diffs = 0.0; - // long double my_diffs = 0.0; - // for (int j = 0; j < data_size; j++) { - // nccl_diffs += abs(nccl_result[j] - ground_truth[j]); - // my_diffs += abs(my_result[j] - ground_truth[j]); - // } - // if (myRank == 0) - // std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size - // << " me: " << my_diffs / data_size << std::endl; - } - - CUDACHECK(cudaFree(result)); - CUDACHECK(cudaFree(self_data_copy)); - CUDACHECK(cudaFree(rank_data)); - CUDACHECK(cudaFree(buffer)); - CUDACHECK(cudaFree(states)); - CUDACHECK(cudaFreeHost(ground_truth)); - CUDACHECK(cudaFreeHost(nccl_result)); - CUDACHECK(cudaFreeHost(my_result)); - CUDACHECK(cudaStreamDestroy(stream)); -} - -int main(int argc, char** argv) { - int nRanks, myRank; - MPICHECK(MPI_Init(&argc, &argv)); - MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank)); - MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks)); - CUDACHECK(cudaSetDevice(myRank)); - ncclUniqueId id; - ncclComm_t comm; - if (myRank == 0) ncclGetUniqueId(&id); - MPICHECK(MPI_Bcast(static_cast(&id), sizeof(id), MPI_BYTE, 0, - MPI_COMM_WORLD)); - NCCLCHECK(ncclCommInitRank(&comm, nRanks, id, myRank)); - - bool performance_test = true; - cudaProfilerStart(); -// Uncomment to scan through different block size configs. -// for (int threads : {256, 512, 1024}) { -// for (int block_limit = 16; block_limit < 112; block_limit += 4) { -// run(myRank, nRanks, comm, threads, block_limit, 1024 * 1024, -// performance_test); -// } -// } -#ifdef USE_ROCM - const int block_limit = 16; -#else - const int block_limit = 36; -#endif - // Scan through different sizes to test performance. - for (int sz = 512; sz <= (8 << 20); sz *= 2) { - run(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test); - } - - cudaProfilerStop(); - MPICHECK(MPI_Finalize()); - return EXIT_SUCCESS; -} \ No newline at end of file diff --git a/csrc/ops.h b/csrc/ops.h index 398ae1016f3..c310bd59ff5 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -41,12 +41,6 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input); void gelu_quick(torch::Tensor& out, torch::Tensor& input); -void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope, - torch::Tensor const& q_pe, - torch::Tensor const& kv_c_and_k_pe_cache, - torch::Tensor const& seq_lens, - torch::Tensor const& page_table, double scale); - void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor const& scale, std::optional const& azp); diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 16e0df0df64..02404a2f517 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2962,70 +2962,6 @@ def qr_max_size() -> int: return torch.ops._C_custom_ar.qr_max_size() -def get_flash_mla_metadata( - cache_seqlens: torch.Tensor, - num_heads_per_head_k: int, - num_heads_k: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - cache_seqlens: (batch_size), dtype torch.int32. - num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k. - num_heads_k: num_heads_k. - - Return: - tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: (batch_size + 1), dtype torch.int32. - """ - return torch.ops._C.get_flash_mla_metadata( - cache_seqlens, num_heads_per_head_k, num_heads_k - ) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: torch.Tensor, - cache_seqlens: torch.Tensor, - head_dim_v: int, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, - softmax_scale: float | None = None, - causal: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. - cache_seqlens: (batch_size), torch.int32. - head_dim_v: Head_dim of v. - tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), torch.int32, return by get_mla_metadata. - num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata. - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim). - causal: bool. Whether to apply causal attention mask. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - out, softmax_lse = torch.ops._C.flash_mla_fwd_kvcache( - q, - k_cache, - None, - head_dim_v, - cache_seqlens, - block_table, - softmax_scale, - causal, - tile_scheduler_metadata, - num_splits, - ) - return out, softmax_lse - - def sm100_cutlass_mla_decode( out: torch.Tensor, lse: torch.Tensor, From 3daea7ceb990bff87e925b2f4b77325af052282f Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 25 Jun 2026 20:03:09 -0600 Subject: [PATCH 014/138] [Bugfix][MRV2] Forward seq_lens_cpu_upper_bound for mamba hybrid models (#46759) Signed-off-by: mgoin --- vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index ced97c4f277..329f008a4e3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -137,6 +137,7 @@ class MambaHybridModelState(DefaultModelState): block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=mamba_attn_metadata, for_cudagraph_capture=for_capture, From 5314665badcb93f798e117aacad8ce02f148cd73 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:29:25 -0500 Subject: [PATCH 015/138] [Model Runner V2][DFlash] Enable dflash attention backend selection (#46770) Signed-off-by: Giancarlo Delfin --- vllm/v1/worker/gpu/spec_decode/dflash/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 01f6923a76a..c4f98e715b9 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -26,7 +26,9 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_vllm_config = replace( vllm_config, attention_config=replace( - vllm_config.attention_config, use_non_causal=not causal + vllm_config.attention_config, + use_non_causal=not causal, + backend=speculative_config.attention_backend, ), ) with set_model_tag("dflash_head"): From 652d962bc9df7e04959e84ce478c3a8d26fe52a7 Mon Sep 17 00:00:00 2001 From: yiheng Date: Fri, 26 Jun 2026 10:30:07 +0800 Subject: [PATCH 016/138] [Model Runner V2][Spec Decode] Reduce TP communication for draft token generation (#46448) Signed-off-by: EanWang211123 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/v1/worker/gpu/spec_decode/speculator.py | 37 ++++++++++++++++++-- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 4fd7cce36b3..b06c9372a95 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -8,6 +8,7 @@ import torch.nn as nn from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import ( @@ -23,6 +24,8 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +logger = init_logger(__name__) + class BaseSpeculator(ABC): @abstractmethod @@ -95,6 +98,9 @@ class DraftModelSpeculator(BaseSpeculator): self.vocab_size = self.draft_model_config.get_vocab_size() self.dtype = vllm_config.model_config.dtype self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel + self.use_local_argmax_reduction = ( + self.speculative_config.use_local_argmax_reduction + ) # DP configuration self.dp_size = vllm_config.parallel_config.data_parallel_size @@ -149,6 +155,7 @@ class DraftModelSpeculator(BaseSpeculator): ) self.model = self.load_draft_model(target_model, target_attn_layer_names) + self._validate_local_argmax_reduction() all_attn_layers = set[str]( get_layers_from_vllm_config( @@ -211,6 +218,31 @@ class DraftModelSpeculator(BaseSpeculator): ) return attn_metadata + def _validate_local_argmax_reduction(self) -> None: + if not self.use_local_argmax_reduction: + return + if self.speculative_config.draft_sample_method == "probabilistic": + raise ValueError( + "use_local_argmax_reduction is not compatible with " + "draft_sample_method='probabilistic'." + ) + if not hasattr(self.model, "get_top_tokens"): + raise ValueError( + "use_local_argmax_reduction is enabled but draft model " + f"{self.model.__class__.__name__} does not implement " + "get_top_tokens()." + ) + logger.info( + "Using local argmax reduction for draft token generation " + "(communication: O(2*tp_size) vs O(vocab_size))." + ) + + def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_local_argmax_reduction: + return self.model.get_top_tokens(hidden_states) + logits = self.model.compute_logits(hidden_states) + return logits.argmax(dim=-1) + def sample_draft( self, hidden_states: torch.Tensor, @@ -221,8 +253,8 @@ class DraftModelSpeculator(BaseSpeculator): draft_step: torch.Tensor, draft_logits: torch.Tensor | None, ) -> torch.Tensor: - logits = self.model.compute_logits(hidden_states) if draft_logits is not None: + logits = self.model.compute_logits(hidden_states) # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise # used for draft and target sampling. return gumbel_sample( @@ -236,8 +268,7 @@ class DraftModelSpeculator(BaseSpeculator): output_processed_logits_col=draft_step, use_fp64=self.use_fp64_gumbel, ) - else: - return logits.argmax(dim=-1) + return self._greedy_sample_draft(hidden_states) def _copy_request_inputs( self, From 02a1f23711c5bdbff81eb8a610dde39e1141d036 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:32:07 -0500 Subject: [PATCH 017/138] [DFlash] Fuse precompute kv per-layer rmsnorms (#46761) Signed-off-by: Giancarlo Delfin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- csrc/libtorch_stable/layernorm_kernels.cu | 33 ++++++--- .../core/test_batched_weight_rms_norm.py | 70 +++++++++++++++++++ vllm/model_executor/models/qwen3_dflash.py | 23 +++--- 3 files changed, 108 insertions(+), 18 deletions(-) create mode 100644 tests/kernels/core/test_batched_weight_rms_norm.py diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index f29734fc265..0c59a09b1b9 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -20,20 +20,27 @@ __global__ void rms_norm_kernel( const int64_t input_stride_d4, // input.stride(-4) const int64_t input_shape_d2, // input.size(-2) const int64_t input_shape_d3, // input.size(-3) - const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight + const scalar_t* __restrict__ weight, // [hidden_size] or + // [num_groups, hidden_size]; + // null if !HasWeight + const int64_t weight_stride, // 0 or weight.stride(0) const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; const scalar_t* input_row; + const scalar_t* weight_row; + int64_t weight_row_off = 0; if constexpr (NUM_DIMS == 2) { // 2D for layernorm normal case [batch_size, hidden] input_row = input + blockIdx.x * input_stride_d2; + weight_row = weight + blockIdx.x * weight_stride; } else if constexpr (NUM_DIMS == 3) { // 3D for q/k norm [batch_size, num_heads, head_size] int batch_idx = blockIdx.x / input_shape_d2; int head_idx = blockIdx.x % input_shape_d2; input_row = input + batch_idx * input_stride_d3 + head_idx * input_stride_d2; + weight_row = weight + batch_idx * weight_stride; } else if constexpr (NUM_DIMS == 4) { // 4D for transformers model_impl qk norm [batch, seq, head, head_dim] int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2); @@ -42,6 +49,7 @@ __global__ void rms_norm_kernel( int head_idx = remaining % input_shape_d2; input_row = input + batch_idx * input_stride_d4 + seq_idx * input_stride_d3 + head_idx * input_stride_d2; + weight_row = weight + batch_idx * weight_stride; } auto vec_op = [&variance](const vec_n_t& vec) { @@ -69,7 +77,7 @@ __global__ void rms_norm_kernel( scalar_t* out_row = out + blockIdx.x * hidden_size; auto* v_in = reinterpret_cast*>(input_row); - auto* v_w = reinterpret_cast*>(weight); + auto* v_w = reinterpret_cast*>(weight_row); auto* v_out = reinterpret_cast*>(out_row); for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) { vec_n_t dst; @@ -211,15 +219,24 @@ fused_add_rms_norm_kernel( void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] torch::stable::Tensor& input, // [..., hidden_size] - std::optional weight, // [hidden_size] - double epsilon) { + std::optional weight, double epsilon) { STD_TORCH_CHECK(out.is_contiguous()); if (input.stride(-1) != 1) { input = torch::stable::contiguous(input); } STD_TORCH_CHECK(input.stride(-1) == 1); + int64_t weight_stride = 0; if (weight.has_value()) { STD_TORCH_CHECK(weight->is_contiguous()); + if (weight->dim() == 1) { + STD_TORCH_CHECK(weight->size(0) == input.size(-1)); + } else if (weight->dim() == 2) { + STD_TORCH_CHECK(weight->size(0) == input.size(0)); + STD_TORCH_CHECK(weight->size(-1) == input.size(-1)); + weight_stride = weight->stride(0); + } else { + STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D"); + } } int hidden_size = input.size(-1); @@ -256,16 +273,16 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] out.mutable_data_ptr(), input.const_data_ptr(), input_stride_d2, input_stride_d3, input_stride_d4, input_shape_d2, - input_shape_d3, weight_ptr, epsilon, num_tokens, - hidden_size); + input_shape_d3, weight_ptr, weight_stride, epsilon, + num_tokens, hidden_size); } else { vllm::rms_norm_kernel <<>>( out.mutable_data_ptr(), input.const_data_ptr(), input_stride_d2, input_stride_d3, input_stride_d4, input_shape_d2, - input_shape_d3, weight_ptr, epsilon, num_tokens, - hidden_size); + input_shape_d3, weight_ptr, /*weight_stride=*/0, epsilon, + num_tokens, hidden_size); } }); }); diff --git a/tests/kernels/core/test_batched_weight_rms_norm.py b/tests/kernels/core/test_batched_weight_rms_norm.py new file mode 100644 index 00000000000..42711fdf09e --- /dev/null +++ b/tests/kernels/core/test_batched_weight_rms_norm.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the batched-weight RMS norm kernel (vllm._custom_ops.rms_norm). + +``rms_norm`` can use the outermost input batch index to select the corresponding +weight row. The result must match that of looping ``rms_norm`` over that dimension. +""" + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="rms_norm requires a CUDA/ROCm device", +) + + +@pytest.mark.parametrize( + "shape", + [ + (28, 17, 128), # 3D: [num_rows, tokens, hidden] + (1, 5, 2, 128), # 4D: single row (edge case) + (28, 13, 8, 128), # 4D: [L, num_ctx, nkv, hd] (DFlash K-norm) + (6, 3, 4, 769), # 4D: non-power-of-two hidden size + ], +) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float]) +@pytest.mark.parametrize("seed", [42]) +@torch.inference_mode() +def test_rms_norm_matches_loop( + shape: tuple[int, ...], dtype: torch.dtype, seed: int +) -> None: + set_random_seed(seed) + torch.set_default_device("cuda") + + num_rows, hidden = shape[0], shape[-1] + eps = 1e-6 + + x = torch.randn(*shape, dtype=dtype) * 0.1 + # Distinct weight per row so that a wrong row index would be caught. + weight = torch.randn(num_rows, hidden, dtype=dtype) * 0.1 + 1.0 + + # Reference batched-weight rms norm. + out_ref = torch.empty_like(x) + for i in range(x.shape[0]): + ops.rms_norm(out_ref[i], x[i], weight[i], eps) + + out = torch.empty_like(x) + ops.rms_norm(out, x, weight, eps) + + # Expect bitwise-identical results. + torch.testing.assert_close(out, out_ref, atol=0, rtol=0) + + +@torch.inference_mode() +def test_rms_norm_validates_shapes() -> None: + torch.set_default_device("cuda") + + x = torch.randn(4, 8, 128, dtype=torch.float) + out = torch.empty_like(x) + # Expect num rows mismatch. + with pytest.raises(RuntimeError): + ops.rms_norm(out, x, torch.randn(3, 128), 1e-6) + # Expect hidden size mismatch. + with pytest.raises(RuntimeError): + ops.rms_norm(out, x, torch.randn(4, 64), 1e-6) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 36c0a357878..8746a15f115 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -309,8 +309,11 @@ class DFlashQwen3Model(nn.Module): else: self._fused_kv_bias = None - # K-norm weights: list of [head_dim] tensors, one per layer. - self._k_norm_weights = [a.k_norm.weight.data for a in layers_attn] + # K-norm weights stacked into one contiguous [num_layers, head_dim] + # tensor so the per-layer K-norm runs as a single grouped kernel. + self._k_norm_weights = torch.stack( + [a.k_norm.weight.data for a in layers_attn], dim=0 + ).contiguous() # RoPE parameters self._rope_head_size = attn0.rotary_emb.head_size @@ -392,15 +395,15 @@ class DFlashQwen3Model(nn.Module): all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous - # --- Per-layer RMSNorm K (3D: [num_ctx, nkv, hd] per layer) --- + # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) --- + # The weight is selected per layer by the outermost (layer) index. all_k_normed = torch.empty_like(all_k) - for i in range(L): - ops.rms_norm( - all_k_normed[i], - all_k[i], - self._k_norm_weights[i], - self._rms_norm_eps, - ) + ops.rms_norm( + all_k_normed, + all_k, + self._k_norm_weights, + self._rms_norm_eps, + ) # --- Fused RoPE across all layers --- # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy). From 552a9dbe59bf2e6a35654440c64c5e52bed90586 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Fri, 26 Jun 2026 04:33:00 +0200 Subject: [PATCH 018/138] [NVFP4][Emulation] Fuse NVFP4 weight dequantization with compute in triton kernel for w13/w2 MOE MLP linears (#44667) Signed-off-by: Felix Marty --- .../quantization/test_nvfp4_emulation.py | 463 ++++++++++++++++++ .../fused_moe/experts/nvfp4_emulation_moe.py | 461 +++++++++++++++-- .../utils/nvfp4_emulation_utils.py | 48 +- 3 files changed, 905 insertions(+), 67 deletions(-) diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index 71072d9e9ff..f5652af6e92 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -1,10 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import cast + import huggingface_hub import pytest import torch from safetensors import safe_open +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, + nvfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import ( + Nvfp4QuantizationEmulationTritonExperts, +) +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts from vllm.model_executor.layers.quantization.utils import ( nvfp4_emulation_utils, ) @@ -12,9 +27,167 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import dequantize_to_dtype, ref_nvfp4_quant_dequant, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kNvfp4Dynamic, + kNvfp4Static, +) from vllm.platforms import current_platform from vllm.triton_utils import triton +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 +else: + + def on_gfx950() -> bool: + return False + + +class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts): + """ + Extension of TritonExperts to support emulated NVFP4 MoE experts. + + It may be used for NVFP4 models when the device does not have + native support for this dtype. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + + # `TritonExperts.apply` expects pre-dequantized weights, + # which we handle in `apply` below. + self.w1_scale_val = self.quant_config.w1_scale + self.w2_scale_val = self.quant_config.w2_scale + + self.quant_config._w1.scale = None + self.quant_config._w2.scale = None + + self.quantization_emulation = True + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return "nvfp4" + + @property + def a1_scale(self) -> torch.Tensor | None: + return self.quant_config.a1_gscale + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) + + 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, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert w1.dtype == torch.uint8 + assert w2.dtype == torch.uint8 + + # Dequantize w1 from packed NVFP4 to fp16/bf16 + w13_global_scale = self.quant_config.g1_alphas + + w1_dequant = dequantize_to_dtype( + tensor_fp4=w1, + tensor_sf=self.w1_scale_val, + global_scale=w13_global_scale, + dtype=hidden_states.dtype, + block_size=16, + swizzle=False, + ) + + # Dequantize w2 from packed NVFP4 to fp16/bf16 + w2_global_scale = self.quant_config.g2_alphas + + w2_dequant = dequantize_to_dtype( + tensor_fp4=w2, + tensor_sf=self.w2_scale_val, + global_scale=w2_global_scale, + dtype=hidden_states.dtype, + block_size=16, + swizzle=False, + ) + + super().apply( + output=output, + hidden_states=hidden_states, + w1=w1_dequant, + w2=w2_dequant, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=None, + a2_scale=self.quant_config.a2_gscale, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + +@pytest.mark.parametrize( + ("config_kwargs", "expected_reason"), + [ + ({"has_bias": True}, "kernel does not support bias"), + ({"is_lora_enabled": True}, "kernel does not support LoRA"), + ], +) +def test_nvfp4_emulation_support_check_rejects_bias_and_lora( + config_kwargs: dict[str, bool], + expected_reason: str, +) -> None: + moe_config = FusedMoEConfig( + num_experts=2, + experts_per_token=1, + hidden_dim=16, + intermediate_size=16, + num_local_experts=2, + num_logical_experts=2, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.TopK, + **config_kwargs, + ) + + supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config( + Nvfp4QuantizationEmulationTritonExperts, + moe_config, + kNvfp4Static, + kNvfp4Dynamic, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == expected_reason + @pytest.mark.skipif( not current_platform.is_cuda_alike(), @@ -306,3 +479,293 @@ def test_triton_nvfp4_quant_dequant( f"min={ref_min:.3f}ms, max={ref_max:.3f}ms" ) print(f" speedup: {speedup:.2f}x") + + +MOE_MODEL_CONFIGS = { + "nvidia/Qwen3-30B-A3B-NVFP4": { + "shards": ["model-00001-of-00004.safetensors"], + "expert_prefix": "model.layers.9.mlp.experts.", + # Position of the expert index in the dot-split key. + "expert_idx_pos": 5, + }, + "nvidia/Kimi-K2.6-NVFP4": { + "shards": [ + "model-00001-of-00060.safetensors", + "model-00002-of-00060.safetensors", + ], + "expert_prefix": "language_model.model.layers.1.mlp.experts.", + "expert_idx_pos": 6, + }, +} + + +def _load_nvfp4_moe_weights( + model_id: str, + tensor_parallel_size: int, + max_experts: int | None = None, +): + """Load and stack NVFP4 MoE weights from checkpoint shards. + + Returns (w1, w1_scale, w1_gscale, w2, w2_scale, w2_gscale, + a1_gscale, a2_gscale, num_experts, hidden_dim, + intermediate_size). + + When max_experts is set, only the first max_experts experts are loaded. + + When tensor_parallel_size > 1, the N dimension of w1 and the K + dimension of w2 are narrowed to the first TP shard (simulating + column-parallel on w1 / row-parallel on w2). + """ + cfg = MOE_MODEL_CONFIGS[model_id] + shards = cast(list[str], cfg["shards"]) + checkpoint_path = huggingface_hub.snapshot_download( + model_id, + allow_patterns=shards, + ) + expert_prefix = cfg["expert_prefix"] + idx_pos = cast(int, cfg["expert_idx_pos"]) + + # Collect all tensors across shards into a flat dict โ€” an expert's + # tensors may be split across multiple shard files. + all_tensors: dict[str, torch.Tensor] = {} + for shard_name in shards: + shard_path = f"{checkpoint_path}/{shard_name}" + with safe_open(shard_path, framework="pt", device="cpu") as f: + for key in f.keys(): # noqa: SIM118 + if key.startswith(expert_prefix): + all_tensors[key] = f.get_tensor(key) + + expert_indices = sorted( + { + int(key.split(".")[idx_pos]) + for key in all_tensors + if key.endswith(".gate_proj.weight") + } + ) + if max_experts is not None: + expert_indices = expert_indices[:max_experts] + num_experts = len(expert_indices) + + gate_weights, up_weights, down_weights = [], [], [] + gate_scales, up_scales, down_scales = [], [], [] + gate_gscales, up_gscales, down_gscales = [], [], [] + a1_scales, a2_scales = [], [] + + for idx in expert_indices: + prefix = f"{expert_prefix}{idx}" + gate_weights.append(all_tensors[f"{prefix}.gate_proj.weight"]) + gate_scales.append(all_tensors[f"{prefix}.gate_proj.weight_scale"]) + gate_gscales.append(all_tensors[f"{prefix}.gate_proj.weight_scale_2"]) + up_weights.append(all_tensors[f"{prefix}.up_proj.weight"]) + up_scales.append(all_tensors[f"{prefix}.up_proj.weight_scale"]) + up_gscales.append(all_tensors[f"{prefix}.up_proj.weight_scale_2"]) + down_weights.append(all_tensors[f"{prefix}.down_proj.weight"]) + down_scales.append(all_tensors[f"{prefix}.down_proj.weight_scale"]) + down_gscales.append(all_tensors[f"{prefix}.down_proj.weight_scale_2"]) + a1_scales.append(all_tensors[f"{prefix}.gate_proj.input_scale"]) + a2_scales.append(all_tensors[f"{prefix}.down_proj.input_scale"]) + + # Stack into MoE format. + # w1 = [E, 2*intermediate, hidden//2] (gate + up concatenated) + w1 = torch.stack( + [torch.cat([g, u], dim=0) for g, u in zip(gate_weights, up_weights)] + ).cuda() + w1_scale = torch.stack( + [torch.cat([g, u], dim=0) for g, u in zip(gate_scales, up_scales)] + ).cuda() + w1_gscale = torch.stack(gate_gscales).cuda() + + # w2 = [E, hidden, intermediate//2] + w2 = torch.stack(down_weights).cuda() + w2_scale = torch.stack(down_scales).cuda() + w2_gscale = torch.stack(down_gscales).cuda() + + a13_scale_raw = torch.stack(a1_scales).cuda() + a2_scale_raw = torch.stack(a2_scales).cuda() + + # Apply EMULATION transforms (matches oracle/nvfp4.py). + nvfp4_emulation_utils.kE2M1ToFloat_handle.val = ( + nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda() + ) + a1_gscale = 1.0 / a13_scale_raw.max().to(torch.float32) + a2_gscale = 1.0 / a2_scale_raw.max().to(torch.float32) + + # โ”€โ”€ Simulate TP sharding โ”€โ”€ + # w1 (gate_up): column-parallel โ†’ shard the N dimension (dim 1). + # w2 (down): row-parallel โ†’ shard the K dimension (dim 2, + # which is the packed K//2 dim). + # Scales follow the same sharding on the corresponding dimension. + tp = tensor_parallel_size + if tp > 1: + n1 = w1.size(1) // tp + w1 = w1[:, :n1, :].contiguous() + w1_scale = w1_scale[:, :n1, :].contiguous() + + k2_packed = w2.size(2) // tp + k2_scale = w2_scale.size(2) // tp + w2 = w2[:, :, :k2_packed].contiguous() + w2_scale = w2_scale[:, :, :k2_scale].contiguous() + + hidden_dim = w1.size(2) * 2 + intermediate_size = w1.size(1) // 2 + + return ( + w1, + w1_scale, + w1_gscale, + w2, + w2_scale, + w2_gscale, + a1_gscale, + a2_gscale, + num_experts, + hidden_dim, + intermediate_size, + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Triton NVFP4 kernel requires CUDA.", +) +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 1024]) +@pytest.mark.parametrize("top_k", [4]) +@pytest.mark.parametrize("model_id", list(MOE_MODEL_CONFIGS.keys())) +@pytest.mark.parametrize( + "tensor_parallel_size", + [pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]], +) +def test_nvfp4_moe_correctness( + num_tokens: int, + top_k: int, + model_id: str, + tensor_parallel_size: int, +) -> None: + """Compare Nvfp4QuantizationEmulationTritonExperts (fused weight dequant + compute) + against the unfused reference Nvfp4QuantizationEmulationTritonExpertsReference. + + Both must produce bit-identical results. + """ + num_test_experts = max(8, top_k) + ( + w1, + w1_scale, + w1_gscale, + w2, + w2_scale, + w2_gscale, + a1_gscale, + a2_gscale, + num_experts, + hidden_dim, + intermediate_size, + ) = _load_nvfp4_moe_weights( + model_id, + tensor_parallel_size, + max_experts=num_test_experts, + ) + + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=top_k, + hidden_dim=hidden_dim, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + + def _make_quant_config(): + return nvfp4_moe_quant_config( + g1_alphas=w1_gscale.clone(), + g2_alphas=w2_gscale.clone(), + a1_gscale=a1_gscale.clone(), + a2_gscale=a2_gscale.clone(), + w1_scale=w1_scale.clone(), + w2_scale=w2_scale.clone(), + ) + + ref_experts = Nvfp4QuantizationEmulationTritonExpertsReference( + moe_config=moe_config, + quant_config=_make_quant_config(), + ) + fused_experts = Nvfp4QuantizationEmulationTritonExperts( + moe_config=moe_config, + quant_config=_make_quant_config(), + ) + + torch.manual_seed(42) + hidden_states = torch.randn( + num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda" + ) + + topk_weights = torch.randn( + num_tokens, top_k, dtype=torch.float32, device="cuda" + ).softmax(dim=-1) + topk_ids = torch.stack( + [torch.randperm(num_experts, device="cuda")[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + N = w1.size(1) # 2 * intermediate + K = hidden_dim + + ws13_size = num_tokens * top_k * max(intermediate_size, K) + ws2_size = num_tokens * top_k * max(N, K) + + workspace13_ref = torch.zeros(ws13_size, dtype=torch.bfloat16, device="cuda") + workspace2_ref = torch.zeros(ws2_size, dtype=torch.bfloat16, device="cuda") + output_ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device="cuda") + + workspace13_fused = torch.zeros_like(workspace13_ref) + workspace2_fused = torch.zeros_like(workspace2_ref) + output_fused = torch.zeros_like(output_ref) + + apply_kwargs = dict( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + # Unfused reference. + ref_experts.apply( + output=output_ref, + workspace13=workspace13_ref, + workspace2=workspace2_ref, + **apply_kwargs, + ) + + # Fused implementation. + fused_experts.apply( + output=output_fused, + workspace13=workspace13_fused, + workspace2=workspace2_fused, + **apply_kwargs, + ) + + # Not strict equality on H100, MI325, MI300 (< 0.1% elements). + # The fused on-the-fly dequant path can lower to a slightly + # different Triton/MMA tiling than the pre-dequantized + # reference; experiments with reference-like tiling/masking + # reduced some diffs were not kept because they regress + # the fused kernel speed. + # Strict equality validated on MI355. + torch.testing.assert_close( + output_fused, + output_ref, + atol=0.0 if on_gfx950() else 0.02, + rtol=0, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py index d7ed53612e0..f93c67a97dd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py @@ -11,6 +11,8 @@ Weights are dequantized on the fly during each forward, we fall back to calling is applied on `a13`, `a2`. """ +from typing import Any + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -21,18 +23,316 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts +from vllm.model_executor.layers.fused_moe.fused_moe import ( + try_get_optimal_moe_config, + write_zeros_to_output, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.model_executor.layers.fused_moe.utils import ( + _resize_cache, + moe_kernel_quantize_input, +) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( - dequantize_to_dtype, + _e2m1_inline, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kNvfp4Dynamic, kNvfp4Static, ) +from vllm.triton_utils import tl, triton logger = init_logger(__name__) +@triton.jit +def fused_moe_nvfp4_emulation_kernel( + a_ptr, + b_ptr, + c_ptr, + b_scale_ptr, + w_global_scale_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + N: tl.constexpr, + K: tl.constexpr, + EM, + num_valid_tokens, + # Strides โ€” A [M, K] + stride_am, + stride_ak, + # Strides โ€” B [E, N, K//2], passed as (expert, K-packed, N) + stride_be, + stride_bk, + stride_bn, + # Strides โ€” C [M, topk, N] + stride_cm, + stride_cn, + # Strides โ€” B_scale [E, N, K//BLOCK], passed as (expert, K-scale, N) + stride_bse, + stride_bsk, + stride_bsn, + block_k_diviable: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + group_size: tl.constexpr, +): + """ + Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM. + + Activations A are BF16 (already QDQ'd externally). + Weights B are packed uint8 NVFP4 [E, N, K//2] โ€” two FP4 values per byte + along the K dimension. + B_scale holds per-block FP8-E4M3 scales [E, N, K // group_size]. + w_global_scale is a per-expert scalar global scale. + + The dequantization formula per element is: + w_float = e2m1_decode(nibble) * (block_scale_fp8 * global_scale) + + Weight loading optimization: each packed byte is loaded exactly once as + a [BLOCK_SIZE_N, BLOCK_SIZE_K // 2] tile (N-major), both nibbles are + extracted, decoded and scaled, then tl.interleave produces the + [BLOCK_SIZE_N, BLOCK_SIZE_K] dequantized tile which is transposed to + [BLOCK_SIZE_K, BLOCK_SIZE_N] for tl.dot. + """ + BLOCK_SIZE_K_PACKED: tl.constexpr = BLOCK_SIZE_K // 2 + + # Map program ids to the block of C it should compute. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # Token / expert setup + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_experts == -1: + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + # Pointer setup + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_k_packed = tl.arange(0, BLOCK_SIZE_K_PACKED) + + # A pointers: [BLOCK_SIZE_M, BLOCK_SIZE_K] + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + # B pointers: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED] โ€” N-major so that + # tl.interleave (which operates on the last dim) produces a + # [BLOCK_SIZE_N, BLOCK_SIZE_K] tile that we transpose for tl.dot. + # Each unique byte is loaded exactly once. + b_ptrs = ( + b_ptr + + off_experts * stride_be + + offs_bn[:, None] * stride_bn + + offs_k_packed[None, :] * stride_bk + ) + + # B_scale pointers: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED] โ€” same + # N-major layout. Each packed byte index covers 2 K elements that + # always fall within the same group (group_size=16, so each group + # spans 8 packed bytes). We can therefore index the scale using + # offs_k_packed directly. + # Note: group_size_packed = group_size // 2 maps packed indices to + # scale indices the same way unpacked indices map via group_size. + group_size_packed: tl.constexpr = group_size // 2 + + # Load per-expert global scale (scalar). + w_global_scale = tl.load(w_global_scale_ptr + off_experts).to(tl.float32) + + # K-loop with FP32 accumulation + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load A tile [BLOCK_SIZE_M, BLOCK_SIZE_K]. + if block_k_diviable: + a = tl.load( + a_ptrs, + mask=token_mask[:, None], + other=0.0, + ) + else: + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + + # Load packed weight tile [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED]. + if block_k_diviable: + raw_bytes = tl.load(b_ptrs) + else: + kp_mask = offs_k_packed[None, :] < (K // 2) - k * BLOCK_SIZE_K_PACKED + raw_bytes = tl.load(b_ptrs, mask=kp_mask, other=0) + + # Extract both nibbles from each byte (each [N, K_packed]). + low_nibble = raw_bytes & 0x0F + high_nibble = (raw_bytes >> 4) & 0x0F + + low_decoded = _e2m1_inline(low_nibble) + high_decoded = _e2m1_inline(high_nibble) + + # Load and apply per-block FP8 scales. + # Scale shape: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED], one scale per + # group_size_packed packed elements. + b_scale_ptrs = ( + b_scale_ptr + + off_experts * stride_bse + + offs_bn[:, None] * stride_bsn + + ((offs_k_packed[None, :] + BLOCK_SIZE_K_PACKED * k) // group_size_packed) + * stride_bsk + ) + if block_k_diviable: + b_scale_raw = tl.load(b_scale_ptrs) + else: + b_scale_raw = tl.load(b_scale_ptrs, mask=kp_mask, other=0.0) + + b_scale = tl.cast(b_scale_raw, tl.float8e4nv, bitcast=True).to(tl.float32) + b_scale = b_scale * w_global_scale + + # Scale both halves with the same per-block scale (the two + # elements packed in one byte always belong to the same group). + low_scaled = low_decoded * b_scale + high_scaled = high_decoded * b_scale + + # Interleave along last dim: [N, K_packed] x2 -> [N, K], + # then transpose to [K, N] for tl.dot. + b = tl.trans(tl.interleave(low_scaled, high_scaled)).to(compute_type) + + accumulator = tl.dot(a, b, acc=accumulator) + + # Advance pointers along K. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K_PACKED * stride_bk + + # Router weight multiplication (in float32 for stability) + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) + accumulator = accumulator * moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + + # Write output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +def invoke_fused_moe_nvfp4_emulation_kernel( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + B_scale: torch.Tensor, + act_global_scale: torch.Tensor, + w_global_scale: torch.Tensor, + topk_weights: torch.Tensor | None, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + mul_routed_weight: bool, + top_k: int, + config: dict[str, Any], + compute_type: tl.dtype, +): + """Launch the fused NVFP4 emulation MoE kernel. + + B has shape [E, N, K_packed] where K_packed = K // 2 (two FP4 per byte). + B_scale has shape [E, N, K // group_size] in FP8-E4M3 (stored as uint8). + w_global_scale has shape [E] (per-expert scalar). + """ + assert B_scale is not None and B_scale.ndim == 3 + + N = B.size(1) + K = A.size(1) + + M = A.size(0) + num_tokens = M * top_k + + EM = sorted_token_ids.size(0) + if A.size(0) < config["BLOCK_SIZE_M"]: + EM = min( + sorted_token_ids.size(0), + A.size(0) * top_k * config["BLOCK_SIZE_M"], + ) + + grid = lambda META: ( + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + fused_moe_nvfp4_emulation_kernel[grid]( + A, + B, + C, + B_scale, + w_global_scale, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + N, + K, + EM, + num_tokens, + A.stride(0), + A.stride(1), + # B is [E, N, K//2]: swap N and K strides so kernel indexes [K, N]. + B.stride(0), + B.stride(2), + B.stride(1), + C.stride(1), + C.stride(2), + # B_scale is [E, N, K//group]: swap N and K strides likewise. + B_scale.stride(0), + B_scale.stride(2), + B_scale.stride(1), + block_k_diviable=K % config["BLOCK_SIZE_K"] == 0, + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=top_k, + compute_type=compute_type, + group_size=16, + BLOCK_SIZE_M=config["BLOCK_SIZE_M"], + BLOCK_SIZE_N=config["BLOCK_SIZE_N"], + BLOCK_SIZE_K=config["BLOCK_SIZE_K"], + GROUP_SIZE_M=config["GROUP_SIZE_M"], + ) + + class Nvfp4QuantizationEmulationTritonExperts(TritonExperts): """ Extension of TritonExperts to support emulated NVFP4 MoE experts. @@ -72,6 +372,31 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts): def expects_unquantized_inputs(self) -> bool: return True + @staticmethod + def supports_lora() -> bool: + return False + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + if moe_config.is_lora_enabled: + return False, "kernel does not support LoRA" + if moe_config.has_bias: + return False, "kernel does not support bias" + + return TritonExperts.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -109,47 +434,109 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts): # w2 shape: [num_experts, hidden_size, intermediate_size//2] assert w1.dtype == torch.uint8 assert w2.dtype == torch.uint8 + assert hidden_states.is_contiguous() + assert hidden_states.dim() == 2 - # Dequantize w1 from packed NVFP4 to fp16/bf16 - w13_global_scale = self.quant_config.g1_alphas + K = hidden_states.size(-1) + assert w1.size(2) * 2 == K, f"Hidden size mismatch: {K} != {w1.size(2) * 2}" - w1_dequant = dequantize_to_dtype( - tensor_fp4=w1, - tensor_sf=self.w1_scale_val, - global_scale=w13_global_scale, - dtype=hidden_states.dtype, - block_size=16, - swizzle=False, + E, num_tokens, N, _, top_k_num = self.moe_problem_size( + hidden_states, w1, w2, topk_ids ) - # Dequantize w2 from packed NVFP4 to fp16/bf16 - w2_global_scale = self.quant_config.g2_alphas + if global_num_experts == -1: + global_num_experts = E - w2_dequant = dequantize_to_dtype( - tensor_fp4=w2, - tensor_sf=self.w2_scale_val, - global_scale=w2_global_scale, - dtype=hidden_states.dtype, - block_size=16, - swizzle=False, + # TODO: There is actually no support for tuning of the underlying triton + # hyperparameters in benchmarks/kernels/benchmark_moe.py, to be added. + config = try_get_optimal_moe_config( + w1.size(), + w2.size(), + top_k_num, + self.quant_config.config_name(hidden_states.dtype), + num_tokens, + block_shape=None, ) - # Activation quantization/dequantization is deferred to - # `moe_kernel_quantize_input` in TritonExperts.apply. - super().apply( - output=output, - hidden_states=hidden_states, - w1=w1_dequant, - w2=w2_dequant, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=activation, - global_num_experts=global_num_experts, - expert_map=expert_map, - a1q_scale=None, - a2_scale=self.quant_config.a2_gscale, - workspace13=workspace13, - workspace2=workspace2, - expert_tokens_meta=expert_tokens_meta, - apply_router_weight_on_input=apply_router_weight_on_input, + if hidden_states.dtype == torch.bfloat16: + compute_type = tl.bfloat16 + elif hidden_states.dtype == torch.float16: + compute_type = tl.float16 + elif hidden_states.dtype == torch.float32: + compute_type = tl.float32 + else: + raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") + + intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N)) + activation_out_dim = self.adjust_N_for_activation(N, activation) + intermediate_cache2 = _resize_cache( + workspace13, (num_tokens * top_k_num, activation_out_dim) ) + intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K)) + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, + config["BLOCK_SIZE_M"], + global_num_experts, + expert_map, + ) + + # Activation NVFP4 QDQ. + hidden_states_qdq, _ = moe_kernel_quantize_input( + A=hidden_states, + A_scale=self.quant_config.a1_gscale, + quant_dtype="nvfp4", + per_act_token_quant=False, + quantization_emulation=True, + ) + + # w13: fused weight dequant + GEMM. + invoke_fused_moe_nvfp4_emulation_kernel( + hidden_states_qdq, + w1, + intermediate_cache1, + self.w1_scale_val, + self.quant_config.a1_gscale, + self.quant_config.g1_alphas, + None, # topk_weights โ€” applied after w2 + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + False, # mul_routed_weight + top_k_num, + config, + compute_type=compute_type, + ) + + self.activation( + activation, intermediate_cache2, intermediate_cache1.view(-1, N) + ) + + # Activation NVFP4 QDQ. + intermediate_cache2_qdq, _ = moe_kernel_quantize_input( + A=intermediate_cache2, + A_scale=self.quant_config.a2_gscale, + quant_dtype="nvfp4", + per_act_token_quant=False, + quantization_emulation=True, + ) + + # w2: fused weight dequant + GEMM. + invoke_fused_moe_nvfp4_emulation_kernel( + intermediate_cache2_qdq, + w2, + intermediate_cache3, + self.w2_scale_val, + self.quant_config.a2_gscale, + self.quant_config.g2_alphas, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + not apply_router_weight_on_input, + 1, + config, + compute_type=compute_type, + ) + + self.moe_sum(intermediate_cache3, output) diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index 39c78a9062b..ad6b272371e 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -23,28 +23,24 @@ kE2M1ToFloat_handle = SimpleNamespace( @triton.jit -def _e2m1_inline(magnitude): - """Inline E2M1 lookup using binary tree - 3 levels instead of 7 sequential. +def _e2m1_inline(nibble): + """Decode an NVFP4 nibble (4 bits: 1 sign + 3 magnitude) to float32. - Maps 3-bit magnitude to float: [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] - Uses bit decomposition for fewer comparisons. + Uses direct IEEE 754 bit construction. + For magnitudes 2-7 the FP32 bit pattern is 0x3F000000 + (mag << 22), + which is a single shift + add + bitcast. Magnitudes 0 (zero) and 1 + (E2M1 subnormal = 0.5) are patched with two tl.where ops. """ - # Bit 2 (MSB): separates 0-3 from 4-7 - # Bit 1: separates within groups - # Bit 0 (LSB): separates within pairs - b2 = (magnitude >> 2) & 1 # 0 for mag 0-3, 1 for mag 4-7 - b1 = (magnitude >> 1) & 1 # middle bit - b0 = magnitude & 1 # LSB + magnitude = nibble & 0x07 + sign = (nibble >> 3) & 1 - # For mag 0-3: [0.0, 0.5, 1.0, 1.5] - low_group = tl.where( - b1 == 1, tl.where(b0 == 1, 1.5, 1.0), tl.where(b0 == 1, 0.5, 0.0) - ) - # For mag 4-7: [2.0, 3.0, 4.0, 6.0] - high_group = tl.where( - b1 == 1, tl.where(b0 == 1, 6.0, 4.0), tl.where(b0 == 1, 3.0, 2.0) - ) - return tl.where(b2 == 1, high_group, low_group) + fp32_bits = 0x3F000000 + (magnitude.to(tl.int32) << 22) + val = fp32_bits.to(tl.float32, bitcast=True) + + val = tl.where(magnitude == 0, 0.0, val) + val = tl.where(magnitude == 1, 0.5, val) + + return tl.where(sign == 1, -val, val) @triton.jit @@ -65,7 +61,7 @@ def _dequantize_nvfp4_kernel( """ BLOCK_PACKED: tl.constexpr = BLOCK_SIZE // 2 - row_idx = tl.program_id(0) + row_idx = tl.program_id(0).to(tl.int64) tile_idx = tl.program_id(1) if has_batch_global_scale: @@ -105,16 +101,8 @@ def _dequantize_nvfp4_kernel( low_nibble = raw_bytes & 0x0F high_nibble = (raw_bytes >> 4) & 0x0F - # Binary tree E2M1 decode - low_mag = low_nibble & 0x07 - low_val = _e2m1_inline(low_mag) - low_sign = (low_nibble >> 3) & 1 - low_result = tl.where(low_sign == 1, -low_val, low_val) * scale_values - - high_mag = high_nibble & 0x07 - high_val = _e2m1_inline(high_mag) - high_sign = (high_nibble >> 3) & 1 - high_result = tl.where(high_sign == 1, -high_val, high_val) * scale_values + low_result = _e2m1_inline(low_nibble) * scale_values + high_result = _e2m1_inline(high_nibble) * scale_values # Interleave for coalesced contiguous store result = tl.interleave(low_result, high_result) From dbc49b6b99d02d6daadc0c8150267e67fbecc446 Mon Sep 17 00:00:00 2001 From: ovidiusm Date: Fri, 26 Jun 2026 04:33:42 +0200 Subject: [PATCH 019/138] [CI][NIXL] Fix NIXL EP import canary for the nixl 1.3.0 wheel and pin nixl==1.3.0 (#45166) Signed-off-by: Ovidiu Mara Signed-off-by: ovidiusm Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- requirements/kv_connectors.txt | 2 +- .../nixl_integration/test_nixl_imports.py | 26 +++---------------- 2 files changed, 4 insertions(+), 24 deletions(-) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index e0d494e9f21..ce920816db3 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,5 +2,5 @@ lmcache >= 0.3.9 # CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0 # until a fixed newer release is verified for runtime images. cupy-cuda13x < 14.1.0 -nixl == 1.2.0 # Required for disaggregated prefill +nixl == 1.3.0 mooncake-transfer-engine >= 0.3.8 diff --git a/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py b/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py index feb03d0d1a9..4422f45847b 100644 --- a/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py +++ b/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py @@ -4,8 +4,6 @@ import importlib import importlib.metadata as metadata -import subprocess -import sys import types import pytest @@ -65,25 +63,7 @@ def test_nixl_and_nixl_ep_imports() -> None: # Exercise the NIXL EP extension used by fused MoE expert parallelism. nixl_ep = importlib.import_module("nixl_ep") print(f"nixl_ep: {nixl_ep.__file__}") + assert nixl_ep.__file__ is not None - nixl_ep_cpp = _import_nixl_ep_cpp(nixl_ep) - assert nixl_ep_cpp.__file__ is not None - extension_file = nixl_ep_cpp.__file__ - print(f"nixl_ep_cpp: {extension_file}") - - completed = subprocess.run( - ["ldd", extension_file], - capture_output=True, - check=False, - text=True, - ) - print(completed.stdout) - if completed.stderr: - print(completed.stderr, file=sys.stderr) - - assert completed.returncode == 0 - if torch.version.cuda is not None: - cuda_major = torch.version.cuda.split(".", maxsplit=1)[0] - expected_cudart = f"libcudart.so.{cuda_major}" - assert expected_cudart in completed.stdout - assert f"{expected_cudart} => not found" not in completed.stdout + # Check that the NIXL EP extension is loaded. + assert nixl_ep.Config is not None From d350fa8dddc6ed1a3a5710473d485ad0d02a6127 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Fri, 26 Jun 2026 11:41:33 +0800 Subject: [PATCH 020/138] [Bugfix][Rust Frontend] Reject min_tokens above max_tokens (#46733) Co-authored-by: Bugen Zhao Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao --- rust/src/chat/src/error.rs | 11 ++++++ rust/src/server/src/error.rs | 41 +++++++++++---------- rust/src/server/src/grpc/mod.rs | 20 +++++++---- rust/src/server/src/grpc/tests.rs | 59 +++++++++++++++++++++++++++++++ rust/src/text/src/error.rs | 23 ++++++++++++ rust/src/text/src/lower.rs | 27 ++++++++++++++ 6 files changed, 153 insertions(+), 28 deletions(-) diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index c472a65601d..da2396c2198 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -74,6 +74,17 @@ pub enum Error { pub type Result = std::result::Result; +impl Error { + /// Whether this error represents invalid user request parameters. + pub fn is_request_validation_error(&self) -> bool { + match self { + Self::PromptTooLong { .. } => true, + Self::Text(error) => error.is_request_validation_error(), + _ => false, + } + } +} + /// Format the available-parser suffix used in user-facing error messages. fn available_parser_hint(available_names: &[String]) -> String { if available_names.is_empty() { diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index ede83748f3d..3eba278267f 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -78,7 +78,7 @@ impl IntoResponse for ApiError { /// the client's fault and map to HTTP 400, mirroring the Python frontend. /// Everything else stays an internal 500. pub fn text_submit_error(context: &'static str, error: vllm_text::Error) -> ApiError { - if is_request_validation_error(&error) { + if error.is_request_validation_error() { return invalid_request!("{error}"); } server_error!("{}: {}", context, error.to_report_string()) @@ -87,27 +87,10 @@ pub fn text_submit_error(context: &'static str, error: vllm_text::Error) -> ApiE /// Like [`text_submit_error`], for the chat pipeline (which both wraps the /// text errors and raises its own prompt-length variant). pub fn chat_submit_error(context: &'static str, error: vllm_chat::Error) -> ApiError { - match &error { - vllm_chat::Error::PromptTooLong { .. } => invalid_request!("{error}"), - vllm_chat::Error::Text(text_error) if is_request_validation_error(text_error) => { - invalid_request!("{error}") - } - _ => server_error!("{}: {}", context, error.to_report_string()), + if error.is_request_validation_error() { + return invalid_request!("{error}"); } -} - -fn is_request_validation_error(error: &vllm_text::Error) -> bool { - matches!( - error, - vllm_text::Error::PromptTooLong { .. } - | vllm_text::Error::EmptyPromptTokenIds { .. } - | vllm_text::Error::Logprobs(_) - | vllm_text::Error::TokenIds(_) - | vllm_text::Error::InvalidThinkingTokenBudget - // An empty tokenized prompt detected later, at request prepare - // time, surfaces through the transparent Llm wrapper. - | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. }) - ) + server_error!("{}: {}", context, error.to_report_string()) } #[cfg(test)] @@ -140,6 +123,22 @@ mod tests { assert!(response.error.message.contains("thinking_token_budget")); } + #[test] + fn min_tokens_above_max_tokens_maps_to_invalid_request() { + let api_error = text_submit_error( + "failed to submit completion request", + vllm_text::Error::MinTokensExceedsMaxTokens { + min_tokens: 5, + max_tokens: 4, + }, + ); + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); + let response = api_error.to_error_response(); + assert_eq!(response.error.error_type, "invalid_request_error"); + assert!(response.error.message.contains("min_tokens=5")); + assert!(response.error.message.contains("max_tokens=4")); + } + #[test] fn chat_wrapped_prompt_too_long_maps_to_invalid_request() { let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong { diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 62ee8607669..1fcb8674fee 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -56,12 +56,9 @@ impl pb::generate_server::Generate for GenerateServiceImpl { info!(%request_id, "grpc generate (unary)"); let stream = self.state.chat.text().generate(text_request).await; - let stream = stream.map_err(|e| Status::internal(e.to_report_string()))?; + let stream = stream.map_err(text_error_to_status)?; - let collected = stream - .collect_output() - .await - .map_err(|e| Status::internal(e.to_report_string()))?; + let collected = stream.collect_output().await.map_err(text_error_to_status)?; // Build the single aggregated response. let prompt_info = convert::to_prompt_info( @@ -104,7 +101,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl { info!(%request_id, "grpc generate (stream)"); let stream = self.state.chat.text().generate(text_request).await; - let stream = stream.map_err(|e| Status::internal(e.to_report_string()))?; + let stream = stream.map_err(text_error_to_status)?; let (tx, rx) = mpsc::channel(32); @@ -112,7 +109,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl { futures::pin_mut!(stream); while let Some(event) = stream.next().await { let response = match event { - Err(e) => Err(Status::internal(e.to_report_string())), + Err(e) => Err(text_error_to_status(e)), Ok(DecodedTextEvent::Start { prompt_token_ids, prompt_logprobs, @@ -155,3 +152,12 @@ impl pb::generate_server::Generate for GenerateServiceImpl { Ok(Response::new(Box::pin(response_stream))) } } + +fn text_error_to_status(error: vllm_text::Error) -> Status { + let message = error.to_report_string(); + if error.is_request_validation_error() { + Status::invalid_argument(message) + } else { + Status::internal(message) + } +} diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 58bf894c920..14156a41046 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -425,6 +425,34 @@ async fn unary_generate_missing_prompt_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { + let (mut client, server_task, _engine_task) = + grpc_test_server(b"engine-grpc-min-above-max", default_stream_output_specs()).await; + + let status = client + .generate(pb::GenerateRequest { + request_id: "test-min-above-max".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 4, + min_new_tokens: 5, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect_err("should fail when min_new_tokens exceeds max_new_tokens"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("min_tokens=5")); + assert!(status.message().contains("max_tokens=4")); + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_yields_incremental_responses() { @@ -513,6 +541,37 @@ async fn streaming_generate_missing_prompt_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn streaming_generate_min_tokens_above_max_tokens_returns_invalid_argument() { + let (mut client, server_task, _engine_task) = grpc_test_server( + b"engine-grpc-stream-min-above-max", + default_stream_output_specs(), + ) + .await; + + let status = client + .generate_stream(pb::GenerateRequest { + request_id: "test-stream-min-above-max".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 4, + min_new_tokens: 5, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect_err("should fail when min_new_tokens exceeds max_new_tokens"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("min_tokens=5")); + assert!(status.message().contains("max_tokens=4")); + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_with_sampling_params() { diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index 2c9e69ca15c..96ddced5841 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -20,6 +20,11 @@ pub enum Error { Logprobs(#[from] LogprobsError), #[error(transparent)] TokenIds(#[from] TokenIdsError), + #[error( + "`min_tokens` must be less than or equal to `max_tokens`, \ + got min_tokens={min_tokens}, max_tokens={max_tokens}" + )] + MinTokensExceedsMaxTokens { min_tokens: u32, max_tokens: u32 }, #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")] InvalidThinkingTokenBudget, #[error("text request stream `{request_id}` closed before terminal output")] @@ -32,6 +37,24 @@ pub enum Error { pub type Result = std::result::Result; +impl Error { + /// Whether this error represents invalid user request parameters. + pub fn is_request_validation_error(&self) -> bool { + match self { + Self::PromptTooLong { .. } + | Self::EmptyPromptTokenIds { .. } + | Self::Logprobs(_) + | Self::TokenIds(_) + | Self::MinTokensExceedsMaxTokens { .. } + | Self::InvalidThinkingTokenBudget + // An empty tokenized prompt detected later, at request prepare + // time, surfaces through the transparent Llm wrapper. + | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) => true, + _ => false, + } + } +} + impl From for Error { fn from(error: vllm_tokenizer::TokenizerError) -> Self { Self::Tokenizer(error.0) diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 38528114e59..6cd18195bb9 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -128,6 +128,12 @@ pub fn lower_sampling_params( prompt_len, )?; let min_tokens = min_tokens.unwrap_or(0); + if min_tokens > max_tokens { + return Err(Error::MinTokensExceedsMaxTokens { + min_tokens, + max_tokens, + }); + } let thinking_token_budget = normalize_thinking_token_budget(thinking_token_budget)?; let frequency_penalty = frequency_penalty.unwrap_or(0.0); let presence_penalty = presence_penalty.unwrap_or(0.0); @@ -414,6 +420,27 @@ mod tests { )); } + #[test] + fn lower_sampling_params_rejects_min_tokens_above_resolved_max_tokens() { + let error = lower_sampling_params_with_limits( + SamplingParams { + max_tokens: Some(4), + min_tokens: Some(5), + ..SamplingParams::default() + }, + sample_sampling_limits(), + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::MinTokensExceedsMaxTokens { + min_tokens: 5, + max_tokens: 4, + } + )); + } + #[test] fn lower_text_request_applies_python_style_eos_hints() { let prepared = lower_text_request( From 1502cf62749cb3ec1cbf00f7ef87c55b2edd2c84 Mon Sep 17 00:00:00 2001 From: Matti4 Date: Fri, 26 Jun 2026 05:45:20 +0200 Subject: [PATCH 021/138] Fix relative allowed local media paths (#45263) --- tests/multimodal/media/test_connector.py | 17 +++++++++++++++++ vllm/multimodal/media/connector.py | 2 +- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index b78d24d189f..bee9d50ac1c 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -152,6 +152,23 @@ async def test_fetch_image_local_files(image_url: str): connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}") +@pytest.mark.asyncio +async def test_fetch_image_local_files_relative_allowed_path(tmp_path, monkeypatch): + media_dir = tmp_path / "media" + media_dir.mkdir() + image_path = media_dir / "image.png" + Image.new("RGB", (1, 1), color=(255, 0, 0)).save(image_path) + + monkeypatch.chdir(tmp_path) + local_connector = MediaConnector(allowed_local_media_path="media") + + image_sync = local_connector.fetch_image(image_path.as_uri()) + image_async = await local_connector.fetch_image_async(image_path.as_uri()) + + assert image_sync.size == (1, 1) + assert not ImageChops.difference(image_sync, image_async).getbbox() + + @pytest.mark.asyncio @pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True) async def test_fetch_image_local_files_with_space_in_name(image_url: str): diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 312239ad3fd..582b6fde565 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -105,7 +105,7 @@ class MediaConnector: self.connection = connection if allowed_local_media_path: - allowed_local_media_path_ = Path(allowed_local_media_path) + allowed_local_media_path_ = Path(allowed_local_media_path).resolve() if not allowed_local_media_path_.exists(): raise ValueError( From e312c5cb25427e76fc3830ab14e7b6bc0963a55c Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Fri, 26 Jun 2026 11:54:03 +0800 Subject: [PATCH 022/138] [Rust Frontend] Make Granite4 string argument scanning incremental (#46507) Signed-off-by: reidliu41 --- rust/src/parser/Cargo.toml | 5 + rust/src/parser/benches/granite4.rs | 75 ++++++++++++ rust/src/parser/src/tool/json/granite4.rs | 75 +++++++++--- rust/src/parser/src/utils.rs | 134 +++++++++++++++++++--- 4 files changed, 256 insertions(+), 33 deletions(-) create mode 100644 rust/src/parser/benches/granite4.rs diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml index 67c74bb5601..09c3d6b5dc1 100644 --- a/rust/src/parser/Cargo.toml +++ b/rust/src/parser/Cargo.toml @@ -74,5 +74,10 @@ name = "gemma4" harness = false required-features = ["test-util"] +[[bench]] +name = "granite4" +harness = false +required-features = ["test-util"] + [lints] workspace = true diff --git a/rust/src/parser/benches/granite4.rs b/rust/src/parser/benches/granite4.rs new file mode 100644 index 00000000000..17ef8671605 --- /dev/null +++ b/rust/src/parser/benches/granite4.rs @@ -0,0 +1,75 @@ +use std::time::Duration; + +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Granite4ToolParser, Tool, ToolParser}; + +mod utils; +use utils::feed_parser; + +const CHUNK_CHARS: usize = 7; +const LONG_ARGUMENT_BYTES: usize = 64 * 1024; + +fn string_args_fixture() -> String { + let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(LONG_ARGUMENT_BYTES)); + let encoded_arguments = serde_json::to_string(&arguments).unwrap(); + format!(r#"{{"name":"f","arguments":{encoded_arguments}}}"#) +} + +fn object_args_fixture() -> String { + format!( + r#"{{"name":"f","arguments":{{"data":"{}"}}}}"#, + "x".repeat(LONG_ARGUMENT_BYTES) + ) +} + +fn parser(tools: &[Tool]) -> Box { + Granite4ToolParser::create(tools).expect("Granite4 parser should initialize") +} + +fn run_stream_group(c: &mut Criterion, name: &str, tools: &[Tool], text: &str) { + let chunks = split_by_chars(text, CHUNK_CHARS); + + let mut group = c.benchmark_group(name); + group.sample_size(50); + group.warm_up_time(Duration::from_millis(300)); + group.measurement_time(Duration::from_secs(2)); + group.throughput(Throughput::Bytes(text.len() as u64)); + + group.bench_function("reuse_parser", |b| { + let mut parser = parser(tools); + b.iter(|| { + let result = feed_parser(&mut *parser, black_box(&chunks)); + debug_assert_eq!(result.0, ""); + debug_assert_eq!(result.1, 1); + black_box(result); + }) + }); + + group.bench_function("create_parser", |b| { + b.iter_batched( + || parser(tools), + |mut parser| { + let result = feed_parser(&mut *parser, black_box(&chunks)); + debug_assert_eq!(result.0, ""); + debug_assert_eq!(result.1, 1); + black_box(result); + }, + BatchSize::SmallInput, + ) + }); + + group.finish(); +} + +fn bench_granite4(c: &mut Criterion) { + let tools = test_tools(); + let string_args = string_args_fixture(); + let object_args = object_args_fixture(); + + run_stream_group(c, "granite4/long_string_arguments", &tools, &string_args); + run_stream_group(c, "granite4/long_object_arguments", &tools, &object_args); +} + +criterion_group!(benches, bench_granite4); +criterion_main!(benches); diff --git a/rust/src/parser/src/tool/json/granite4.rs b/rust/src/parser/src/tool/json/granite4.rs index fe1cf190225..4989578011d 100644 --- a/rust/src/parser/src/tool/json/granite4.rs +++ b/rust/src/parser/src/tool/json/granite4.rs @@ -9,7 +9,8 @@ use super::{ tool_call_header_event, }; use crate::tool::utils::{ - JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object, + JsonObjectScanState, JsonStringScanState, decode_json_str, parse_buffered_event, safe_text_len, + take_json_object, take_json_string, }; use crate::tool::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; @@ -22,14 +23,20 @@ enum Granite4Mode { Header, /// Parsing the arguments value: /// `None` until the first byte decides object vs string; - /// `Some` while streaming an object value. + /// `Some` while streaming the selected value shape. Args { - json_scan: Option, + args_scan: Option, }, /// Arguments done; consume the object's closing `}` and ``. Close, } +#[derive(Debug, Clone, PartialEq, Eq)] +enum Granite4ArgsScan { + Object(JsonObjectScanState), + String(JsonStringScanState), +} + #[derive(Debug, Clone, PartialEq, Eq)] enum Granite4Event { Text { @@ -92,7 +99,7 @@ impl Granite4ToolParser { let tool_index = self.emitted_tool_count; self.emitted_tool_count += 1; self.active_tool_index = Some(tool_index); - self.mode = Granite4Mode::Args { json_scan: None }; + self.mode = Granite4Mode::Args { args_scan: None }; output.push_call(ToolCallDelta { tool_index, name: Some(function_name), @@ -187,7 +194,7 @@ fn parse_next_granite4_event( match mode { Granite4Mode::Text => text_event(input), Granite4Mode::Header => header_event(input), - Granite4Mode::Args { json_scan } => args_event(input, json_scan), + Granite4Mode::Args { args_scan } => args_event(input, args_scan), Granite4Mode::Close => close_event(input), } } @@ -237,14 +244,19 @@ fn header_event(input: &mut JsonToolInput<'_>) -> ModalResult { /// once seen whole and unescaped. fn args_event( input: &mut JsonToolInput<'_>, - json_scan: &mut Option, + args_scan: &mut Option, ) -> ModalResult { - if let Some(scan) = json_scan { - let len = take_json_object(input, scan)?; - return Ok(Granite4Event::ObjectArgsDelta { - len, - complete: scan.complete(), - }); + if let Some(scan) = args_scan { + return match scan { + Granite4ArgsScan::Object(scan) => { + let len = take_json_object(input, scan)?; + Ok(Granite4Event::ObjectArgsDelta { + len, + complete: scan.complete(), + }) + } + Granite4ArgsScan::String(scan) => string_args_event(input, scan), + }; } match peek(any).parse_next(input)? { @@ -252,12 +264,16 @@ fn args_event( let mut scan = JsonObjectScanState::default(); let len = take_json_object(input, &mut scan)?; let complete = scan.complete(); - *json_scan = Some(scan); + *args_scan = Some(Granite4ArgsScan::Object(scan)); Ok(Granite4Event::ObjectArgsDelta { len, complete }) } - '"' => Ok(Granite4Event::StringArgs { - decoded: json_str(input)?, - }), + '"' => { + *args_scan = Some(Granite4ArgsScan::String(JsonStringScanState::default())); + let Some(Granite4ArgsScan::String(scan)) = args_scan else { + unreachable!("Granite4 string scan state was just initialized") + }; + string_args_event(input, scan) + } _ => { let mut error = ContextError::new(); error.push(StrContext::Label("Granite4 arguments")); @@ -266,6 +282,17 @@ fn args_event( } } +fn string_args_event( + input: &mut JsonToolInput<'_>, + scan: &mut JsonStringScanState, +) -> ModalResult { + let text = **input; + let len = take_json_string(input, scan)?; + Ok(Granite4Event::StringArgs { + decoded: decode_json_str(&text[..len])?, + }) +} + /// Parse the tool-call object's closing `}` and the `` end marker. fn close_event(input: &mut JsonToolInput<'_>) -> ModalResult { seq!(_: ws0, _: literal("}"), _: ws0, _: literal(TOOL_CALL_END)) @@ -420,6 +447,22 @@ mod tests { assert_eq!(output.calls()[0].arguments, r#"{"a":1}"#); } + #[test] + fn granite4_long_string_args_stream_without_reparse() { + let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(64 * 1024)); + let encoded_arguments = serde_json::to_string(&arguments).unwrap(); + let input = + format!(r#"{{"name":"f","arguments":{encoded_arguments}}}"#); + let chunks = split_by_chars(&input, 7); + let mut parser = Granite4ToolParser::new(&test_tools()); + + let output = collect_stream(&mut parser, &chunks); + + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("f")); + assert_eq!(output.calls()[0].arguments, arguments); + } + #[test] fn granite4_streaming_handles_marker_and_json_whitespace() { // Granite spaces the markers (` {โ€ฆ} `) and the JSON diff --git a/rust/src/parser/src/utils.rs b/rust/src/parser/src/utils.rs index 70255215393..bd8f6f48e9e 100644 --- a/rust/src/parser/src/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -290,8 +290,22 @@ pub fn take_json_object( Ok(text.len()) } -/// Parse a JSON string literal. -pub fn json_str(input: &mut Partial<&str>) -> ModalResult { +/// Streaming lexical state for a JSON string literal. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct JsonStringScanState { + scanned_len: usize, + escape: bool, +} + +/// Parse a raw JSON string literal, resuming from the last scanned byte. +/// +/// The returned length covers the quoted JSON string. This only scans for the +/// string boundary; callers that need the decoded value should pass the raw +/// slice to [`decode_json_str`]. +pub fn take_json_string( + input: &mut Partial<&str>, + state: &mut JsonStringScanState, +) -> ModalResult { let text = **input; if text.is_empty() { return incomplete(); @@ -305,37 +319,58 @@ pub fn json_str(input: &mut Partial<&str>) -> ModalResult { )); } - let mut escape = false; - let mut index = 1; + let mut index = if state.scanned_len == 0 { + 1 + } else if state.scanned_len <= bytes.len() { + state.scanned_len + } else { + return incomplete(); + }; + while index < bytes.len() { let byte = bytes[index]; index += 1; - if escape { - escape = false; + if state.escape { + state.escape = false; continue; } match byte { - b'\\' => escape = true, + b'\\' => state.escape = true, b'"' => { - let raw = &text[..index]; - let value = serde_json::from_str::(raw).map_err(|_| { - json_scan_error( - "JSON string", - StrContextValue::Description("valid JSON string"), - ) - })?; input.next_slice(index); - return Ok(value); + return Ok(index); } _ => {} } } + state.scanned_len = text.len(); incomplete() } +/// Parse a JSON string literal. +pub fn json_str(input: &mut Partial<&str>) -> ModalResult { + let text = **input; + let checkpoint = input.checkpoint(); + let mut state = JsonStringScanState::default(); + let len = take_json_string(input, &mut state)?; + decode_json_str(&text[..len]).inspect_err(|_| { + input.reset(&checkpoint); + }) +} + +/// Decode a complete JSON string literal. +pub fn decode_json_str(raw: &str) -> ModalResult { + serde_json::from_str::(raw).map_err(|_| { + json_scan_error( + "JSON string", + StrContextValue::Description("valid JSON string"), + ) + }) +} + fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode { let mut error = ContextError::new(); error.push(StrContext::Label(label)); @@ -388,8 +423,8 @@ mod tests { use winnow::stream::{Offset, Partial, Stream}; use super::{ - JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len, - safe_text_len_mul, take_json_object, take_until_marker, + JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len, + safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker, }; #[test] @@ -714,6 +749,71 @@ mod tests { .assert_eq(&error.to_string()); } + #[test] + fn take_json_string_consumes_complete_string() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new(r#""say_\"hi\u0021" rest"#); + let checkpoint = input.checkpoint(); + + let len = take_json_string(&mut input, &mut state).unwrap(); + + assert_eq!(len, r#""say_\"hi\u0021""#.len()); + assert_eq!(input.offset_from(&checkpoint), len); + assert_eq!(*input, " rest"); + } + + #[test] + fn take_json_string_resumes_after_incomplete_input() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new(r#""{\"data\":\"partial"#); + let checkpoint = input.checkpoint(); + + let error = take_json_string(&mut input, &mut state).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + assert_eq!(input.offset_from(&checkpoint), 0); + assert_eq!(state.scanned_len, r#""{\"data\":\"partial"#.len()); + + let mut input = Partial::new(r#""{\"data\":\"partial string\"}" tail"#); + let len = take_json_string(&mut input, &mut state).unwrap(); + + assert_eq!(len, r#""{\"data\":\"partial string\"}""#.len()); + assert_eq!(*input, " tail"); + } + + #[test] + fn take_json_string_tracks_escape_across_chunks() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new(r#""abc\"#); + + let error = take_json_string(&mut input, &mut state).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + assert!(state.escape); + + let mut input = Partial::new(r#""abc\"def" tail"#); + let len = take_json_string(&mut input, &mut state).unwrap(); + + assert_eq!(len, r#""abc\"def""#.len()); + assert_eq!(*input, " tail"); + } + + #[test] + fn take_json_string_rejects_non_string_start() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new("42"); + + let error = take_json_string(&mut input, &mut state).unwrap_err(); + + let ErrMode::Cut(error) = error else { + panic!("expected cut error"); + }; + expect![[r#" + invalid JSON string + expected `"`"#]] + .assert_eq(&error.to_string()); + } + #[test] fn json_str_decodes_escaped_content() { let mut input = Partial::new(r#""say_\"hi\u0021" rest"#); From 1a4984520ed06560db66ae21bbb11362fe82d0bd Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Fri, 26 Jun 2026 00:05:12 -0500 Subject: [PATCH 023/138] [Hardware][AMD][CI] Fix AMD CI image build (#46792) Signed-off-by: Matthew Wong --- csrc/libtorch_stable/layernorm_kernels.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 0c59a09b1b9..de0a103fdf2 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -29,7 +29,6 @@ __global__ void rms_norm_kernel( float variance = 0.0f; const scalar_t* input_row; const scalar_t* weight_row; - int64_t weight_row_off = 0; if constexpr (NUM_DIMS == 2) { // 2D for layernorm normal case [batch_size, hidden] input_row = input + blockIdx.x * input_stride_d2; From 5b33041746b9b9ab45bdbd9b42cdd5d19357879a Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 25 Jun 2026 22:10:36 -0700 Subject: [PATCH 024/138] [ModelRunner V2] Fix whisper test (#46773) --- vllm/v1/worker/gpu/mm/encoder_cache.py | 3 +++ vllm/v1/worker/gpu/model_states/mm_pruning.py | 6 +++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu/mm/encoder_cache.py b/vllm/v1/worker/gpu/mm/encoder_cache.py index 1fcbe642994..065df2975c4 100644 --- a/vllm/v1/worker/gpu/mm/encoder_cache.py +++ b/vllm/v1/worker/gpu/mm/encoder_cache.py @@ -12,6 +12,9 @@ class EncoderCache: # MM hash -> encoder outputs self.encoder_outputs: dict[str, torch.Tensor] = {} + def __len__(self) -> int: + return len(self.encoder_outputs) + def add_request( self, req_id: str, mm_features: list[MultiModalFeatureSpec] ) -> None: diff --git a/vllm/v1/worker/gpu/model_states/mm_pruning.py b/vllm/v1/worker/gpu/model_states/mm_pruning.py index e1eb0987929..781baa6d15f 100644 --- a/vllm/v1/worker/gpu/model_states/mm_pruning.py +++ b/vllm/v1/worker/gpu/model_states/mm_pruning.py @@ -121,10 +121,10 @@ def maybe_create_mm_pruner( ) -> MultiModalPruner | None: """Create a MultiModalPruner if the model prunes embeddings and uses M-RoPE.""" if ( - not rope_state + rope_state is None or not rope_state.has_delta - or not encoder_cache - or not model_config.multimodal_config + or encoder_cache is None + or model_config.multimodal_config is None or not model_config.multimodal_config.is_multimodal_pruning_enabled() or not supports_multimodal_pruning(model) ): From 915e99ec6701b64af2c37c2172b151bf6b04ebbc Mon Sep 17 00:00:00 2001 From: peizhang56 Date: Thu, 25 Jun 2026 22:37:47 -0700 Subject: [PATCH 025/138] [ROCm][Bugfix] Fix HIP fork re-init in multimodal offline examples (#46741) Signed-off-by: pei.zhang --- .buildkite/test-amd.yaml | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 3a6568c2e0a..598940e3e3e 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1374,8 +1374,11 @@ steps: - python3 basic/offline_inference/score.py # Multi-modal models - python3 generate/multimodal/audio_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 + # These two examples import transformers before vllm, so on ROCm the HIP context + # is initialized in the parent before vllm sets this guard, poisoning fork. Set it + # inline to keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). + - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_offline.py --seed 0 + - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 # Pooling models - python3 pooling/embed/vision_embedding_offline.py --seed 0 @@ -1817,7 +1820,10 @@ steps: - pytest -v -s tests/models/test_transformers.py - pytest -v -s tests/models/multimodal/test_mapping.py - python3 examples/basic/offline_inference/chat.py - - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl + # This example imports transformers before vllm, so on ROCm the HIP context is + # initialized in the parent before vllm sets this guard, poisoning fork. Set it + # inline to keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). + - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper #------------------------------------------------------- mi300 ยท quantization --------------------------------------------------------# @@ -2893,8 +2899,11 @@ steps: - python3 basic/offline_inference/score.py # Multi-modal models - python3 generate/multimodal/audio_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_offline.py --seed 0 - - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 + # These two examples import transformers before vllm, so on ROCm the HIP context + # is initialized in the parent before vllm sets this guard, poisoning fork. Set it + # inline to keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). + - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_offline.py --seed 0 + - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 # Pooling models - python3 pooling/embed/vision_embedding_offline.py --seed 0 From 35a49fcfc2295d04fb252c1080f8d7bd9d888e25 Mon Sep 17 00:00:00 2001 From: peizhang56 Date: Thu, 25 Jun 2026 22:38:26 -0700 Subject: [PATCH 026/138] [CI][Bugfix] Spawn engine in mm cache sleep test to fix ROCm HIP error (#46749) Signed-off-by: pei.zhang Co-authored-by: Claude --- tests/multimodal/test_cache.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py index 30df1d831a0..bf297946f46 100644 --- a/tests/multimodal/test_cache.py +++ b/tests/multimodal/test_cache.py @@ -32,6 +32,8 @@ from vllm.multimodal.inputs import ( from vllm.multimodal.processing import PromptInsertion from vllm.utils.mem_constants import GiB_bytes, MiB_bytes +from ..utils import create_new_process_for_each_test + pytestmark = pytest.mark.cpu_test @@ -559,6 +561,7 @@ _SLEEP_VISION_PROMPT = ( ) +@create_new_process_for_each_test() @pytest.mark.skipif( not torch.cuda.is_available(), reason="sleep mode regression requires a CUDA GPU", From c7645bce044be01c3c30ceead99388f01aa6d496 Mon Sep 17 00:00:00 2001 From: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Date: Fri, 26 Jun 2026 14:02:10 +0800 Subject: [PATCH 027/138] Remove grok model arch from vllm (#46706) Signed-off-by: Xianbao QIAN --- docs/configuration/optimization.md | 2 +- docs/models/supported_models.md | 5 - tests/models/language/generation/test_grok.py | 43 - tests/models/registry.py | 4 - tests/tokenizers_/test_basic.py | 5 - vllm/config/model.py | 2 - vllm/model_executor/models/grok1.py | 792 ------------------ vllm/model_executor/models/registry.py | 4 +- .../model_executor/models/transformers/moe.py | 3 - vllm/renderers/grok2.py | 90 -- vllm/renderers/registry.py | 1 - vllm/tokenizers/grok2.py | 452 ---------- vllm/tokenizers/registry.py | 1 - 13 files changed, 3 insertions(+), 1401 deletions(-) delete mode 100644 tests/models/language/generation/test_grok.py delete mode 100644 vllm/model_executor/models/grok1.py delete mode 100644 vllm/renderers/grok2.py delete mode 100644 vllm/tokenizers/grok2.py diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index 32e7726cb15..c6d64b25035 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -297,7 +297,7 @@ The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't, vLLM raises a clear `ImportError` at tokenizer load. The override applies to any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`, `deepseek_v32`, `deepseek_v4`, โ€ฆ). Models that don't use the HF -fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag. +fast tokenizer (`mistral`, `kimi_audio`) ignore the flag. Tokenizer-bound workloads โ€” long shared prefixes, bursty short prompts, batch detokenization โ€” see the largest wins. If your bottleneck is GPU diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 294c0c6b3f2..b0e0e3ce9c4 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -414,8 +414,6 @@ th { | `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | โœ…๏ธŽ | โœ…๏ธŽ | | `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | โœ…๏ธŽ | โœ…๏ธŽ | | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | โœ…๏ธŽ | โœ…๏ธŽ | -| `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | โœ…๏ธŽ | โœ…๏ธŽ | -| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | โœ…๏ธŽ | โœ…๏ธŽ | | `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | โœ…๏ธŽ | โœ…๏ธŽ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | โœ…๏ธŽ | โœ…๏ธŽ | @@ -488,9 +486,6 @@ th { | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | โœ…๏ธŽ | โœ…๏ธŽ | | `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | | -!!! note - Grok2 requires `tokenizer.tok.json` with `tiktoken` installed. You can optionally override MoE router renormalization with `moe_router_renormalize`. - Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it! | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | diff --git a/tests/models/language/generation/test_grok.py b/tests/models/language/generation/test_grok.py deleted file mode 100644 index a2f1e8b4413..00000000000 --- a/tests/models/language/generation/test_grok.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import pytest - -from ...utils import dummy_hf_overrides - -MODELS = ["xai-org/grok-2"] - - -def _grok2_dummy_overrides(hf_config): - hf_config = dummy_hf_overrides(hf_config, model_arch="Grok1ForCausalLM") - text_config = hf_config.get_text_config() - text_config.update( - { - "hidden_size": 256, - "intermediate_size": 512, - "moe_intermediate_size": 256, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 64, - } - ) - return hf_config - - -@pytest.mark.parametrize("model", MODELS) -def test_dummy_generate(vllm_runner, monkeypatch, model: str) -> None: - with monkeypatch.context() as m: - m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - with vllm_runner( - model, - load_format="dummy", - max_model_len=128, - hf_overrides=_grok2_dummy_overrides, - enforce_eager=True, - ) as llm: - prompt = "Hello from Grok-2" - tokenizer = llm.get_llm().get_tokenizer() - prompt_len = len(tokenizer.encode(prompt)) - outputs = llm.generate_greedy([prompt], max_tokens=1) - output_ids, output_str = outputs[0] - assert len(output_ids) > prompt_len - assert output_str is not None diff --git a/tests/models/registry.py b/tests/models/registry.py index bd2cba46b67..be271ea0777 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -319,10 +319,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "GraniteMoeSharedForCausalLM": _HfExamplesInfo( "ibm-research/moe-7b-1b-active-shared-experts" ), - "Grok1ModelForCausalLM": _HfExamplesInfo( - "hpcai-tech/grok-1", trust_remote_code=True - ), - "Grok1ForCausalLM": _HfExamplesInfo("xai-org/grok-2", trust_remote_code=True), "HrmTextForCausalLM": _HfExamplesInfo( "sapientinc/HRM-Text-1B", min_transformers_version="5.9.0", diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index c3549e2c942..fc4da3f8fec 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -9,7 +9,6 @@ from transformers import ( ) from vllm.tokenizers import TokenizerLike, get_tokenizer -from vllm.tokenizers.grok2 import Grok2Tokenizer from vllm.tokenizers.hf import HfTokenizer from vllm.tokenizers.mistral import MistralTokenizer @@ -35,10 +34,6 @@ def test_tokenizer_like_protocol(): assert isinstance(tokenizer, MistralTokenizer) _assert_tokenizer_like(tokenizer) - tokenizer = get_tokenizer("xai-org/grok-2", tokenizer_mode="grok2") - assert isinstance(tokenizer, Grok2Tokenizer) - _assert_tokenizer_like(tokenizer) - tokenizer = get_tokenizer("deepseek-ai/DeepSeek-V3", tokenizer_mode="deepseek_v32") assert isinstance(tokenizer, HfTokenizer) diff --git a/vllm/config/model.py b/vllm/config/model.py index 245af557df0..c7736f985df 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -601,8 +601,6 @@ class ModelConfig: if self.tokenizer_mode == "auto": if self.model_impl == "terratorch": self.tokenizer_mode = "terratorch" - elif arch == "Grok1ForCausalLM": - self.tokenizer_mode = "grok2" elif arch == "MoonshotKimiaForCausalLM": self.tokenizer_mode = "kimi_audio" elif arch == "DeepseekV32ForCausalLM": diff --git a/vllm/model_executor/models/grok1.py b/vllm/model_executor/models/grok1.py deleted file mode 100644 index 3fc3d1a2d2c..00000000000 --- a/vllm/model_executor/models/grok1.py +++ /dev/null @@ -1,792 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from -# https://github.com/ROCm/vllm/blob/cea7419f151cc50293a05b7fac8547f8f887c9f6/vllm/model_executor/models/grok1.py -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Inference-only Grok (Grok1/Grok2) model.""" - -import math -from collections.abc import Iterable -from itertools import islice -from typing import Any - -import torch -import torch.nn.functional as F -from torch import nn - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.logger import init_logger -from vllm.model_executor.layers.activation import GeluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - fused_moe_make_expert_params_mapping, -) -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - ReplicatedLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) -from vllm.sequence import IntermediateTensors - -from .interfaces import SupportsLoRA, SupportsPP -from .utils import ( - AutoWeightsLoader, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - -# Default Grok1-specific constants, overridden by config values if present -DEFAULT_ATTN_OUTPUT_MULTIPLIER = 0.08838834764831845 -DEFAULT_OUTPUT_MULTIPLIER_SCALE = 0.5773502691896257 -DEFAULT_EMBEDDING_MULTIPLIER_SCALE = 78.38367176906169 -DEFAULT_ROUTER_LOGIT_SOFTCAP = 30.0 - -logger = init_logger(__name__) - - -def _get_num_experts(config) -> int: - return getattr(config, "num_experts", getattr(config, "num_local_experts", 8)) - - -def _get_moe_intermediate_size(config) -> int: - return getattr(config, "moe_intermediate_size", config.intermediate_size) - - -def _get_grok_version(config) -> str: - """Detect Grok version from HF config using multiple heuristics.""" - # Check for Grok2-specific attributes (both for robust detection) - has_residual_moe = getattr(config, "residual_moe", False) - has_moe_intermediate_size = hasattr(config, "moe_intermediate_size") - - if has_residual_moe or has_moe_intermediate_size: - return "grok2" - - return "grok1" # Default to Grok1 - - -def _get_rope_parameters(config) -> dict[str, Any] | None: - rope_parameters = getattr(config, "rope_parameters", None) - if rope_parameters is None: - rope_type = getattr(config, "rope_type", None) - if rope_type is None: - return None - rope_parameters = {"rope_type": rope_type} - rope_theta = getattr(config, "rope_theta", None) - if rope_theta is not None: - rope_parameters["rope_theta"] = rope_theta - scaling_factor = getattr(config, "scaling_factor", None) - if scaling_factor is not None: - rope_parameters["factor"] = scaling_factor - for name in ( - "original_max_position_embeddings", - "extrapolation_factor", - "attn_factor", - "beta_fast", - "beta_slow", - ): - value = getattr(config, name, None) - if value is not None: - rope_parameters[name] = value - - if rope_parameters.get("rope_type") == "original": - rope_parameters = dict(rope_parameters) - rope_parameters["rope_type"] = "default" - return rope_parameters - - -def _get_moe_renormalize(config) -> bool: - explicit_value = getattr( - config, "moe_router_renormalize", getattr(config, "moe_renormalize", None) - ) - if explicit_value is not None: - return bool(explicit_value) - return not getattr(config, "residual_moe", False) - - -class Grok1MLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - input_size=hidden_size, - output_sizes=[intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - input_size=intermediate_size, - output_size=hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - self.act_fn = GeluAndMul() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x, _ = self.gate_up_proj(x) - x = self.act_fn(x) - x, _ = self.down_proj(x) - return x - - -class Grok1MoE(nn.Module): - """A tensor-parallel MoE implementation for Grok1 that shards each expert - across all ranks. - - Each expert's weights are sharded across all ranks and a fused MoE - kernel is used for the forward pass, and finally we reduce the outputs - across ranks. - """ - - def __init__( - self, - num_experts: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - router_logit_soft_cap: float = 0.0, - params_dtype: torch.dtype | None = None, - quant_config: QuantizationConfig | None = None, - tp_size: int | None = None, - renormalize: bool = False, - prefix: str = "", - ): - super().__init__() - self.hidden_size = hidden_size - - # Gate always runs at half / full precision for now. - self.gate = ReplicatedLinear( - hidden_size, - num_experts, - bias=False, - params_dtype=params_dtype, - quant_config=None, - prefix=f"{prefix}.gate", - ) - - self.experts = FusedMoE( - num_experts=num_experts, - top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - params_dtype=params_dtype, - renormalize=renormalize, - quant_config=quant_config, - tp_size=tp_size, - activation="gelu", - prefix=f"{prefix}.experts", - ) - self.router_logit_soft_cap = router_logit_soft_cap - - 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_states = hidden_states.view(-1, self.hidden_size) - # router_logits: (num_tokens, n_experts) - router_logits, _ = self.gate(hidden_states) - if self.router_logit_soft_cap > 0: - router_logits = self.router_logit_soft_cap * F.tanh( - router_logits / self.router_logit_soft_cap - ) - final_hidden_states = self.experts(hidden_states, router_logits) - return final_hidden_states.view(orig_shape) - - -class Grok1Attention(nn.Module): - def __init__( - self, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - max_position: int = 4096 * 32, - rope_parameters: dict[str, Any] | None = None, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - config=None, # Added config parameter - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.config = config # Store config reference - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = hidden_size // self.total_num_heads - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=rope_parameters, - is_neox_style=True, - ) - - attn_logits_soft_cap = max(getattr(config, "attn_logit_softcapping", 30.0), 0.0) - attn_logit_softcapping_method = getattr( - config, "attn_logit_softcapping_method", None - ) - if attn_logit_softcapping_method not in (None, "tanh"): - logger.warning_once( - "Grok attention logit softcapping method '%s' is not " - "supported; falling back to default behavior.", - attn_logit_softcapping_method, - ) - - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - logits_soft_cap=attn_logits_soft_cap, - prefix=f"{prefix}.attn", - ) - self.attn_multiplier = ( - getattr(self.config, "attn_output_multiplier", 1.0) if self.config else 1.0 - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - output *= self.attn_multiplier - return output - - -class Grok1DecoderLayer(nn.Module): - def __init__( - self, - config, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - # Check for fp8 quantization - self.use_fp8 = False - if quant_config is not None: - self.use_fp8 = getattr(quant_config, "is_fp8_w8a8", lambda: False)() - if not self.use_fp8 and hasattr(quant_config, "is_fp8"): - self.use_fp8 = quant_config.is_fp8 - - self.attn = Grok1Attention( - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - max_position=config.max_position_embeddings, - num_kv_heads=config.num_key_value_heads, - rope_parameters=_get_rope_parameters(config), - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - config=config, - ) # Pass config to Grok1Attention - - num_experts = _get_num_experts(config) - num_experts_per_tok = getattr(config, "num_experts_per_tok", 2) - moe_intermediate_size = _get_moe_intermediate_size(config) - moe_renormalize = _get_moe_renormalize(config) - - self.moe_block = Grok1MoE( - num_experts=num_experts, - top_k=num_experts_per_tok, - hidden_size=config.hidden_size, - intermediate_size=moe_intermediate_size, - router_logit_soft_cap=max( - getattr( - config, - "router_logit_softcapping", - DEFAULT_ROUTER_LOGIT_SOFTCAP, - ), - 0.0, - ), - quant_config=quant_config, - renormalize=moe_renormalize, - prefix=f"{prefix}.moe_block", - ) - self.residual_moe = getattr(config, "residual_moe", False) - self.residual_moe_scale = 1.0 / math.sqrt(2.0) - - self.pre_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.mlp = None - if self.residual_moe: - self.mlp = Grok1MLP( - hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.pre_attn_norm(hidden_states) - else: - hidden_states, residual = self.pre_attn_norm(hidden_states, residual) - - hidden_states = self.attn( - positions=positions, - hidden_states=hidden_states, - ) - - # Post attention normalization - hidden_states = self.post_attn_norm(hidden_states) - - # MoE block with normalization - hidden_states, residual = self.pre_moe_norm(hidden_states, residual) - if self.residual_moe: - assert self.mlp is not None - hidden_states = ( - self.moe_block(hidden_states) + self.mlp(hidden_states) - ) * self.residual_moe_scale - else: - hidden_states = self.moe_block(hidden_states) - hidden_states = self.post_moe_norm(hidden_states) - - return hidden_states, residual - - -@support_torch_compile -class Grok1Model(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - ckpt_gate_proj_name: str = "linear", - ckpt_down_proj_name: str = "linear_1", - ckpt_up_proj_name: str = "linear_v", - weight_name_remapping: dict[str, str] | None = None, - ): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - self.config = config - self.quant_config = quant_config - - # Store expert naming for weight loading - self.ckpt_gate_proj_name = ckpt_gate_proj_name - self.ckpt_down_proj_name = ckpt_down_proj_name - self.ckpt_up_proj_name = ckpt_up_proj_name - self.weight_name_remapping = weight_name_remapping or {} - - self.vocab_size = config.vocab_size - - self.embedding_multiplier_scale = getattr( - config, "embedding_multiplier_scale", DEFAULT_EMBEDDING_MULTIPLIER_SCALE - ) - - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - quant_config=quant_config, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: Grok1DecoderLayer( - config, cache_config, quant_config=quant_config, prefix=prefix - ), - prefix=f"{prefix}.layers", - ) - - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - hidden_states = self.embed_tokens(input_ids) - hidden_states = hidden_states * self.embedding_multiplier_scale - return hidden_states - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer(positions, hidden_states, residual) - - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - # Map expert parameter names to standard names - num_experts = _get_num_experts(self.config) - return fused_moe_make_expert_params_mapping( - self, - ckpt_gate_proj_name=self.ckpt_gate_proj_name, - ckpt_down_proj_name=self.ckpt_down_proj_name, - ckpt_up_proj_name=self.ckpt_up_proj_name, - num_experts=num_experts, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("mlp.gate_up_proj", "mlp.gate_proj", 0), - ("mlp.gate_up_proj", "mlp.up_proj", 1), - ] - - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - expert_params_mapping = self.get_expert_mapping() - for name, loaded_weight in weights: - # Apply version-specific weight name remapping - for old_pattern, new_pattern in self.weight_name_remapping.items(): - if old_pattern in name: - name = name.replace(old_pattern, new_pattern) - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if ( - name.endswith(".bias") or name.endswith("_bias") - ) and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - if name.endswith("scale"): - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - if ( - name.endswith(".bias") or name.endswith("_bias") - ) and name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader( - param, - loaded_weight, - name, - shard_id=shard_id, - expert_id=expert_id, - ) - break - else: - # Skip loading extra bias for GPTQ models. - if ( - name.endswith(".bias") or name.endswith("_bias") - ) and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - # Handle Grok1-specific norm.scale naming - if "norm.scale" in name: - name = name.replace("scale", "weight") - - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -class GrokBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP): - """Base class for Grok models with shared logic.""" - - fall_back_to_pt_during_load = False - - # Subclasses should override these - packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - } - - # Expert weight naming - subclasses override these - ckpt_gate_proj_name: str = "linear" - ckpt_down_proj_name: str = "linear_1" - ckpt_up_proj_name: str = "linear_v" - - def get_weight_name_remapping(self) -> dict[str, str]: - """Return weight name remapping for this version. Override in subclasses.""" - return {} - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - self.quant_config = quant_config - - self.model = Grok1Model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "model"), - ckpt_gate_proj_name=self.ckpt_gate_proj_name, - ckpt_down_proj_name=self.ckpt_down_proj_name, - ckpt_up_proj_name=self.ckpt_up_proj_name, - weight_name_remapping=self.get_weight_name_remapping(), - ) - - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - - if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight - - self.output_multiplier_scale = getattr( - config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE - ) - self.logits_processor = LogitsProcessor( - config.vocab_size, - scale=self.output_multiplier_scale, - soft_cap=getattr(config, "final_logit_softcapping", None), - ) - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Skip lm_head when tie_word_embeddings is True - skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else None - - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, - ) - return loader.load_weights(weights) - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - return self.model.get_expert_mapping() - - -class Grok1ForCausalLM(GrokBaseForCausalLM): - """Grok1-specific implementation.""" - - # Grok1 expert weight naming - ckpt_gate_proj_name = "linear" - ckpt_down_proj_name = "linear_1" - ckpt_up_proj_name = "linear_v" - - def get_weight_name_remapping(self) -> dict[str, str]: - # Grok1 uses standard naming, no remapping needed - return {} - - -class Grok2ForCausalLM(GrokBaseForCausalLM): - """Grok2-specific implementation.""" - - # Grok2 has additional packed modules for MLP - packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], - } - - # Grok2 expert weight naming - ckpt_gate_proj_name = "w1" - ckpt_down_proj_name = "w2" - ckpt_up_proj_name = "w3" - - def get_weight_name_remapping(self) -> dict[str, str]: - # Grok2 checkpoint uses different naming conventions - return { - ".self_attn.": ".attn.", - ".block_sparse_moe.": ".moe_block.", - } - - -# Version dispatch mapping -_GROK_VERSIONS: dict[str, type[GrokBaseForCausalLM]] = { - "grok1": Grok1ForCausalLM, - "grok2": Grok2ForCausalLM, -} - - -class GrokForCausalLM(GrokBaseForCausalLM): - """Factory class that dispatches to version-specific implementation.""" - - def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""): - config = vllm_config.model_config.hf_config - version = _get_grok_version(config) - - instance_cls = _GROK_VERSIONS.get(version) - if instance_cls is None: - raise ValueError(f"Unsupported Grok version: {version}") - - # Merge class attributes for LoRA/quantization compatibility - cls.packed_modules_mapping = dict(cls.packed_modules_mapping) - cls.packed_modules_mapping.update(instance_cls.packed_modules_mapping) - - return instance_cls(vllm_config=vllm_config, prefix=prefix) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 0a90d9f9c28..1f9e3a24fe4 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -124,8 +124,6 @@ _TEXT_GENERATION_MODELS = { "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), "GritLM": ("gritlm", "GritLM"), - "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"), - "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"), "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"), "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"), "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"), @@ -731,6 +729,8 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "BaichuanForCausalLM": "0.23.0", "AquilaModel": "0.24.0", "AquilaForCausalLM": "0.24.0", + "Grok1ModelForCausalLM": "0.24.0", + "Grok1ForCausalLM": "0.24.0", } _OOT_SUPPORTED_MODELS = { diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 60e39b330f0..372c5b1ec12 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -180,7 +180,6 @@ class MoEMixin(MixtureOfExperts): # (ckpt_gate_proj_name, ckpt_down_proj_name, ckpt_up_proj_name) ("gate_proj", "down_proj", "up_proj"), # Most common MoE style ("w1", "w2", "w3"), # Granite, Mixtral, Phi MoE style - ("linear", "linear_1", "linear_v"), # Grok1 style ] num_experts = self.model_config.get_num_experts() num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts @@ -238,8 +237,6 @@ class MoEMixin(MixtureOfExperts): wrapped_arch = self.config.architectures[0].lower() if "gptoss" in wrapped_arch: activation = "swigluoai" - elif "grok1" in wrapped_arch: - activation = "gelu" # Expert mapping for `AutoWeightsLoader` expert_mapping = self.get_expert_mapping() diff --git a/vllm/renderers/grok2.py b/vllm/renderers/grok2.py deleted file mode 100644 index 665d9a98e94..00000000000 --- a/vllm/renderers/grok2.py +++ /dev/null @@ -1,90 +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 ( - ChatCompletionMessageParam, - ConversationMessage, - parse_chat_messages, - parse_chat_messages_async, -) -from vllm.logger import init_logger -from vllm.tokenizers.grok2 import Grok2Tokenizer -from vllm.utils.async_utils import make_async - -from .base import BaseRenderer -from .inputs import DictPrompt -from .inputs.preprocess import parse_dec_only_prompt -from .params import ChatParams - -logger = init_logger(__name__) - - -class Grok2Renderer(BaseRenderer[Grok2Tokenizer]): - def __init__( - self, - config: VllmConfig, - tokenizer: Grok2Tokenizer | None, - ) -> None: - super().__init__(config, tokenizer) - - self._apply_chat_template_async = make_async( - self._apply_chat_template, executor=self._executor - ) - - def _apply_chat_template(self, *args, **kwargs): - return self.get_tokenizer().apply_chat_template(*args, **kwargs) - - def render_messages( - self, - messages: list[ChatCompletionMessageParam], - params: ChatParams, - ) -> tuple[list[ConversationMessage], DictPrompt]: - conversation, mm_data, mm_uuids = parse_chat_messages( - messages, - self.model_config, - content_format="string", - media_io_kwargs=params.media_io_kwargs, - mm_processor_kwargs=params.mm_processor_kwargs, - ) - - prompt_raw = self._apply_chat_template( - conversation=conversation, - messages=messages, - **params.get_apply_chat_template_kwargs(), - ) - - prompt = parse_dec_only_prompt(prompt_raw) - if mm_data is not None: - prompt["multi_modal_data"] = mm_data - if mm_uuids is not None: - prompt["multi_modal_uuids"] = mm_uuids - - return conversation, prompt - - async def render_messages_async( - self, - messages: list[ChatCompletionMessageParam], - params: ChatParams, - ) -> tuple[list[ConversationMessage], DictPrompt]: - conversation, mm_data, mm_uuids = await parse_chat_messages_async( - messages, - self.model_config, - content_format="string", - media_io_kwargs=params.media_io_kwargs, - mm_processor_kwargs=params.mm_processor_kwargs, - ) - - prompt_raw = await self._apply_chat_template_async( - conversation=conversation, - messages=messages, - **params.get_apply_chat_template_kwargs(), - ) - - prompt = parse_dec_only_prompt(prompt_raw) - if mm_data is not None: - prompt["multi_modal_data"] = mm_data - if mm_uuids is not None: - prompt["multi_modal_uuids"] = mm_uuids - - return conversation, prompt diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py index a6da9ec5017..098a58e8edc 100644 --- a/vllm/renderers/registry.py +++ b/vllm/renderers/registry.py @@ -22,7 +22,6 @@ logger = init_logger(__name__) _VLLM_RENDERERS = { "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"), - "grok2": ("grok2", "Grok2Renderer"), "hf": ("hf", "HfRenderer"), "kimi_audio": ("hf", "HfRenderer"), "mistral": ("mistral", "MistralRenderer"), diff --git a/vllm/tokenizers/grok2.py b/vllm/tokenizers/grok2.py deleted file mode 100644 index 612af537408..00000000000 --- a/vllm/tokenizers/grok2.py +++ /dev/null @@ -1,452 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tokenizer for Grok-2 .tok.json format.""" - -import functools -import json -from collections.abc import Collection, Sequence, Set -from pathlib import Path -from typing import Any, Literal, overload - -from huggingface_hub.utils import ( - EntryNotFoundError, - HfHubHTTPError, - RepositoryNotFoundError, - RevisionNotFoundError, -) -from transformers import BatchEncoding -from transformers.utils import chat_template_utils as hf_chat_utils - -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam -from vllm.logger import init_logger -from vllm.transformers_utils.repo_utils import hf_api - -from .protocol import TokenizerLike - -logger = init_logger(__name__) - -PAD = "<|pad|>" -EOS = "<|eos|>" -SEP = "<|separator|>" -RESERVED_TOKEN_TEXTS = [f"<|reserved_{i}|>" for i in range(3, 128)] -CONTROL_TOKEN_TEXTS = [f"<|control{i}|>" for i in range(1, 705)] -DEFAULT_SPECIAL_TOKENS = [PAD, SEP, EOS] -DEFAULT_CONTROL_TOKENS = {"pad": PAD, "sep": SEP, "eos": EOS} -DEFAULT_CHAT_TEMPLATE = ( - "{% for message in messages %}" - "{% if message['role'] == 'user' %}" - "{{ 'Human: ' + message['content'].strip() + '<|separator|>\\n\\n' }}" - "{% elif message['role'] == 'system' %}" - "{{ 'System: ' + message['content'].strip() + '<|separator|>\\n\\n' }}" - "{% elif message['role'] == 'assistant' %}" - "{{ 'Assistant: ' + message['content'] + '<|separator|>\\n\\n' }}" - "{% endif %}" - "{% endfor %}" - "{% if add_generation_prompt %}" - "{{ 'Assistant:' }}" - "{% endif %}" -) - -# Default + separate each single digit. -PAT_STR_B = ( - r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|""" - r""" ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""" -) - - -def _maybe_load_tokenizer_config( - model_path: Path, - *, - repo_id: str | None, - revision: str | None, - download_dir: str | None, -) -> dict[str, Any]: - config_path = model_path / "tokenizer_config.json" - if config_path.is_file(): - with config_path.open("r", encoding="utf-8") as f: - return json.load(f) - - if repo_id is None: - return {} - - try: - config_file = hf_api().hf_hub_download( - repo_id=repo_id, - filename="tokenizer_config.json", - revision=revision, - cache_dir=download_dir, - ) - except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError): - # If the repo, revision, or file does not exist, fall back silently. - return {} - except HfHubHTTPError as exc: - logger.warning( - "Failed to download tokenizer_config.json from %s. " - "This may be due to a network or authentication issue. " - "The default chat template will be used. Error: %s", - repo_id, - exc, - ) - return {} - - try: - with Path(config_file).open("r", encoding="utf-8") as f: - return json.load(f) - except json.JSONDecodeError as exc: - logger.warning( - "Failed to parse tokenizer_config.json. " - "The default chat template will be used. Error: %s", - exc, - ) - return {} - except OSError as exc: - logger.warning( - "Failed to open tokenizer_config.json. " - "The default chat template will be used. Error: %s", - exc, - ) - return {} - - -def _load_tiktoken_encoding( - vocab_file: Path, -) -> tuple[Any, dict[str, int]]: - try: - import tiktoken - except ImportError as exc: - raise ImportError("Grok-2 tokenizer requires the `tiktoken` package.") from exc - - with vocab_file.open("rb") as f: - xtok_dict = json.load(f) - - mergeable_ranks = { - bytes(item["bytes"]): item["token"] - for item in xtok_dict.get("regular_tokens", []) - } - special_tokens = { - bytes(item["bytes"]).decode("utf-8", errors="replace"): item["token"] - for item in xtok_dict.get("special_tokens", []) - } - - if xtok_dict.get("word_split") == "V1": - pat_str = PAT_STR_B - else: - raise ValueError(f"Unknown word_split: {xtok_dict.get('word_split')!r}") - - pat_str = xtok_dict.get("pat_str", pat_str) - - kwargs = { - "name": str(vocab_file), - "pat_str": pat_str, - "mergeable_ranks": mergeable_ranks, - "special_tokens": special_tokens, - } - - if "vocab_size" in xtok_dict: - kwargs["explicit_n_vocab"] = xtok_dict["vocab_size"] - - tokenizer = tiktoken.Encoding(**kwargs) - - default_allowed_special: set[str] | None = None - if "default_allowed_special" in xtok_dict: - default_allowed_special = { - bytes(bytes_list).decode("utf-8", errors="replace") - for bytes_list in xtok_dict["default_allowed_special"] - } - - tokenizer._default_allowed_special = default_allowed_special or set() - tokenizer._control_tokens = DEFAULT_CONTROL_TOKENS - - def encode_patched( - self, - text: str, - *, - allowed_special: Literal["all"] | Set[str] = set(), - disallowed_special: Literal["all"] | Collection[str] = "all", - ) -> list[int]: - del disallowed_special - if isinstance(allowed_special, set): - allowed_special |= self._default_allowed_special - return tiktoken.Encoding.encode( - self, - text, - allowed_special=allowed_special, - disallowed_special=(), - ) - - tokenizer.encode = functools.partial(encode_patched, tokenizer) - tokenizer._default_allowed_special |= set(DEFAULT_CONTROL_TOKENS.values()) - tokenizer._default_allowed_special |= set( - CONTROL_TOKEN_TEXTS + RESERVED_TOKEN_TEXTS - ) - - return tokenizer, special_tokens - - -class Grok2Tokenizer(TokenizerLike): - @classmethod - def from_pretrained( - cls, - path_or_repo_id: str | Path, - *args, - trust_remote_code: bool = False, - revision: str | None = None, - download_dir: str | None = None, - **kwargs, - ) -> "Grok2Tokenizer": - if args: - logger.debug_once("Ignoring extra positional args for Grok2Tokenizer.") - - path = Path(path_or_repo_id) - if path.is_file(): - vocab_file = path - model_path = path.parent - repo_id = None - elif path.is_dir(): - vocab_file = path / "tokenizer.tok.json" - model_path = path - repo_id = None - else: - vocab_file = Path( - hf_api().hf_hub_download( - repo_id=str(path_or_repo_id), - filename="tokenizer.tok.json", - revision=revision, - cache_dir=download_dir, - ) - ) - model_path = vocab_file.parent - repo_id = str(path_or_repo_id) - - if not vocab_file.is_file(): - raise FileNotFoundError(f"tokenizer.tok.json not found at {vocab_file}.") - - config = _maybe_load_tokenizer_config( - model_path, - repo_id=repo_id, - revision=revision, - download_dir=download_dir, - ) - - return cls( - vocab_file=vocab_file, - name_or_path=str(path_or_repo_id), - truncation_side=kwargs.get("truncation_side", "left"), - chat_template=config.get("chat_template"), - init_kwargs=config, - ) - - def __init__( - self, - *, - vocab_file: Path, - name_or_path: str, - truncation_side: str, - chat_template: str | None, - init_kwargs: dict[str, Any] | None = None, - ) -> None: - super().__init__() - self.name_or_path = name_or_path - self._truncation_side = truncation_side - self.init_kwargs = init_kwargs or {} - self._chat_template = chat_template or DEFAULT_CHAT_TEMPLATE - - self._tokenizer, self._special_tokens = _load_tiktoken_encoding(vocab_file) - - self._token_to_id: dict[str, int] = {} - self._id_to_token: dict[int, str] = {} - for token, token_id in self._tokenizer._mergeable_ranks.items(): - token_str = token.decode("utf-8", errors="replace") - self._token_to_id[token_str] = token_id - self._id_to_token[token_id] = token_str - - for token, token_id in self._special_tokens.items(): - self._token_to_id[token] = token_id - self._id_to_token[token_id] = token - - bos_token_id = self._special_tokens.get(SEP) - if bos_token_id is None: - bos_token_id = self._special_tokens.get(PAD) - if bos_token_id is None: - bos_token_id = self._special_tokens.get(EOS) - if bos_token_id is None: - bos_token_id = 0 - self._bos_token_id = bos_token_id - - self._eos_token_id = self._special_tokens.get(EOS, self._bos_token_id) - self._pad_token_id = self._special_tokens.get(PAD, self._eos_token_id) - self._unk_token_id = self._pad_token_id - - self._max_chars_per_token = max(len(tok) for tok in self._token_to_id) - - def num_special_tokens_to_add(self) -> int: - return 0 - - @property - def all_special_tokens(self) -> list[str]: - return list(self._special_tokens.keys()) - - @property - def all_special_ids(self) -> list[int]: - return list(self._special_tokens.values()) - - @property - def bos_token_id(self) -> int: - return self._bos_token_id - - @property - def eos_token_id(self) -> int: - return self._eos_token_id - - @property - def pad_token_id(self) -> int: - return self._pad_token_id - - @property - def is_fast(self) -> bool: - return False - - @property - def vocab_size(self) -> int: - return self._tokenizer.n_vocab - - @property - def max_token_id(self) -> int: - return self._tokenizer.n_vocab - 1 - - @property - def max_chars_per_token(self) -> int: - return self._max_chars_per_token - - @property - def truncation_side(self) -> str: - return self._truncation_side - - def get_vocab(self) -> dict[str, int]: - return dict(self._token_to_id) - - def get_added_vocab(self) -> dict[str, int]: - return dict(self._special_tokens) - - def _maybe_truncate(self, tokens: list[int], max_length: int | None) -> list[int]: - if max_length is None or len(tokens) <= max_length: - return tokens - if self.truncation_side == "left": - return tokens[-max_length:] - return tokens[:max_length] - - def encode( - self, - text: str, - truncation: bool | None = None, - max_length: int | None = None, - add_special_tokens: bool = True, - ) -> list[int]: - del add_special_tokens - tokens = self._tokenizer.encode(text) - if truncation: - tokens = self._maybe_truncate(tokens, max_length) - return tokens - - def decode( - self, ids: Sequence[int] | int, skip_special_tokens: bool = False - ) -> str: - if isinstance(ids, int): - ids = [ids] - if skip_special_tokens: - ids = [ - token_id - for token_id in ids - if token_id not in self._special_tokens.values() - ] - return self._tokenizer.decode(ids) - - @overload - def convert_tokens_to_ids(self, tokens: str) -> int: ... - - @overload - def convert_tokens_to_ids(self, tokens: list[str]) -> list[int]: ... - - def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]: - if isinstance(tokens, str): - return self._token_to_id.get(tokens, self._unk_token_id) - return [self._token_to_id.get(token, self._unk_token_id) for token in tokens] - - def convert_ids_to_tokens( - self, ids: Sequence[int], skip_special_tokens: bool = False - ) -> list[str]: - tokens = [] - for token_id in ids: - if skip_special_tokens and token_id in self._special_tokens.values(): - continue - tokens.append(self._id_to_token.get(token_id, "<|unk|>")) - return tokens - - def convert_tokens_to_string(self, tokens: list[str]) -> str: - token_ids = self.convert_tokens_to_ids(tokens) - return self.decode(token_ids, skip_special_tokens=False) - - def __call__( - self, - text: str | list[str], - text_pair: str | None = None, - add_special_tokens: bool = True, - truncation: bool = False, - max_length: int | None = None, - ) -> BatchEncoding: - if text_pair is not None: - raise NotImplementedError("text_pair is not supported for Grok2Tokenizer.") - - if isinstance(text, list): - input_ids_batch: list[list[int]] = [ - self.encode( - item, - truncation=truncation, - max_length=max_length, - add_special_tokens=add_special_tokens, - ) - for item in text - ] - attention_mask_batch = [[1] * len(ids) for ids in input_ids_batch] - return BatchEncoding( - {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} - ) - - input_ids = self.encode( - text, - truncation=truncation, - max_length=max_length, - add_special_tokens=add_special_tokens, - ) - attention_mask = [1] * len(input_ids) - return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask}) - - def get_chat_template( - self, chat_template: str | None, tools: list[dict[str, Any]] | None = None - ) -> str | None: - del tools - return chat_template or self._chat_template - - def apply_chat_template( - self, - messages: list[ChatCompletionMessageParam], - tools: list[dict[str, Any]] | None = None, - chat_template: str | None = None, - tokenize: bool = False, - **kwargs, - ) -> str | list[int]: - template = self.get_chat_template(chat_template, tools=tools) - if template is None: - raise ValueError( - "No chat template available. Provide `chat_template` explicitly." - ) - kwargs["return_dict"] = False - prompt = hf_chat_utils.apply_chat_template( - conversation=messages, - chat_template=template, - tools=tools, - **kwargs, - ) - if tokenize: - return self.encode(prompt, add_special_tokens=False) - return prompt diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index d928da3306e..eb7f8b0cf0d 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -36,7 +36,6 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"} _VLLM_TOKENIZERS = { "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"), - "grok2": ("grok2", "Grok2Tokenizer"), "hf": ("hf", "CachedHfTokenizer"), "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"), "mistral": ("mistral", "MistralTokenizer"), From 63e161f2965e77b2c3ffcd159ce45b2157a21b43 Mon Sep 17 00:00:00 2001 From: Joe Rowell Date: Fri, 26 Jun 2026 08:05:16 +0200 Subject: [PATCH 028/138] [Bugfix][Tool Parser] PoolsideV1: fix string whitespace and required named tool choice (#46486) Signed-off-by: Joe Rowell --- .../test_poolside_v1_tool_parser.py | 217 ++++++++++++++++++ vllm/tool_parsers/poolside_v1_tool_parser.py | 41 +++- 2 files changed, 247 insertions(+), 11 deletions(-) create mode 100644 tests/tool_parsers/test_poolside_v1_tool_parser.py diff --git a/tests/tool_parsers/test_poolside_v1_tool_parser.py b/tests/tool_parsers/test_poolside_v1_tool_parser.py new file mode 100644 index 00000000000..68342e2763b --- /dev/null +++ b/tests/tool_parsers/test_poolside_v1_tool_parser.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for ``PoolsideV1ToolParser``. + +Covers two bugs: + +1. ``adjust_request`` did not skip the forced ``structured_outputs`` JSON + for ``required``/named tool choice. These models emit XML tool calls + (``......``) per the chat + template, so guided JSON decoding conflicts with the format: the call + leaks as content with empty ``tool_calls``. ``adjust_request`` now skips + the constraint for both ChatCompletion (``ChatCompletionNamedToolChoice``) + and Responses (``ToolChoiceFunction``) named choices. + +2. ``extract_tool_calls`` stripped string-typed argument values, corrupting + content whose whitespace is significant (e.g. code/file bodies losing + leading indent and trailing newline). String values are now kept verbatim; + only non-string types are stripped/deserialized. +""" + +from __future__ import annotations + +import json +from typing import Any + +from openai.types.responses.tool_param import FunctionToolParam + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.tool_parsers.poolside_v1_tool_parser import PoolsideV1ToolParser + + +def _write_file_tool() -> dict[str, Any]: + """Tool with a string arg (``content``) and a non-string arg (``mode``).""" + return { + "type": "function", + "function": { + "name": "write_file", + "description": "Write content to a file", + "parameters": { + "type": "object", + "properties": { + "content": {"type": "string"}, + "mode": {"type": "integer"}, + }, + "required": ["content"], + }, + }, + } + + +def _responses_write_file_tool() -> FunctionToolParam: + return FunctionToolParam( + type="function", + name="write_file", + description="Write content to a file", + parameters={ + "type": "object", + "properties": { + "content": {"type": "string"}, + "mode": {"type": "integer"}, + }, + "required": ["content"], + }, + strict=True, + ) + + +def _build_chat_request(*, tool_choice: str | dict[str, Any]) -> ChatCompletionRequest: + return ChatCompletionRequest.model_validate( + { + "model": "poolside-test", + "messages": [{"role": "user", "content": "write the file"}], + "tools": [_write_file_tool()], + "tool_choice": tool_choice, + } + ) + + +def _build_responses_request(*, tool_choice: str | dict[str, Any]) -> ResponsesRequest: + return ResponsesRequest( + model="poolside-test", + input=[{"role": "user", "content": "write the file"}], + tools=[_responses_write_file_tool()], + tool_choice=tool_choice, + stream=True, + max_output_tokens=200, + ) + + +class _StubTokenizer: + """Minimal tokenizer stub to satisfy ``PoolsideV1ToolParser.__init__``.""" + + def get_vocab(self) -> dict[str, int]: + return {"": 151_657, "": 151_658} + + +def _make_parser(request: ChatCompletionRequest) -> PoolsideV1ToolParser: + return PoolsideV1ToolParser(_StubTokenizer(), tools=request.tools) + + +# --------------------------------------------------------------------------- +# Bug 1: required/named must skip forced structured_outputs (#39870 pattern) +# --------------------------------------------------------------------------- + + +def test_required_skips_structured_outputs_chatcompletion() -> None: + request = _build_chat_request(tool_choice="required") + _make_parser(request).adjust_request(request) + + assert request.structured_outputs is None + assert request.skip_special_tokens is False + + +def test_named_skips_structured_outputs_chatcompletion() -> None: + request = _build_chat_request( + tool_choice={"type": "function", "function": {"name": "write_file"}} + ) + _make_parser(request).adjust_request(request) + + assert request.structured_outputs is None + assert request.skip_special_tokens is False + + +def test_required_skips_structured_outputs_responses() -> None: + request = _build_responses_request(tool_choice="required") + PoolsideV1ToolParser(_StubTokenizer()).adjust_request(request) + + assert request.text is None + assert request.skip_special_tokens is False + + +def test_named_skips_structured_outputs_responses() -> None: + # Responses-API named choice parses to ToolChoiceFunction, a different + # type than the ChatCompletion named choice; both must be handled. + request = _build_responses_request( + tool_choice={"type": "function", "name": "write_file"} + ) + PoolsideV1ToolParser(_StubTokenizer()).adjust_request(request) + + assert request.text is None + assert request.skip_special_tokens is False + + +def test_auto_still_keeps_special_tokens() -> None: + request = _build_chat_request(tool_choice="auto") + _make_parser(request).adjust_request(request) + + assert request.skip_special_tokens is False + + +# --------------------------------------------------------------------------- +# Bug 2: string arg whitespace must be preserved (#42026 pattern) +# --------------------------------------------------------------------------- + + +def test_string_arg_preserves_whitespace() -> None: + request = _build_chat_request(tool_choice="auto") + parser = _make_parser(request) + + content = " def f():\n return 1\n" + model_output = ( + "write_file\n" + "content\n" + f"{content}\n" + "" + ) + + result = parser.extract_tool_calls(model_output, request) + + assert result.tools_called + args = json.loads(result.tool_calls[0].function.arguments) + # Leading indent and trailing newline must survive verbatim. + assert args["content"] == content + + +def test_non_string_arg_still_deserialized() -> None: + request = _build_chat_request(tool_choice="auto") + parser = _make_parser(request) + + model_output = ( + "write_file\n" + "content\n" + "hi\n" + "mode\n" + " 420 \n" + "" + ) + + result = parser.extract_tool_calls(model_output, request) + + assert result.tools_called + args = json.loads(result.tool_calls[0].function.arguments) + assert args["content"] == "hi" + # Non-string value is stripped and parsed to its native type. + assert args["mode"] == 420 + + +def test_responses_extract_tool_calls_with_flat_tools() -> None: + # required/named Responses calls route into extract_tool_calls with flat + # FunctionTool (.name); _is_string_type must not raise. + request = _build_responses_request(tool_choice="required") + parser = PoolsideV1ToolParser(_StubTokenizer(), tools=request.tools) + + content = " x = 1\n" + model_output = ( + "write_file\n" + "content\n" + f"{content}\n" + "" + ) + + result = parser.extract_tool_calls(model_output, request) + + assert result.tools_called + args = json.loads(result.tool_calls[0].function.arguments) + assert args["content"] == content diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index e515e1ce637..f5d996176b8 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -17,10 +17,12 @@ from typing import Any import partial_json_parser.core.complete import regex as re +from openai.types.responses import ToolChoiceFunction from partial_json_parser.core.options import Allow 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 ( @@ -53,6 +55,8 @@ class PoolsideV1ToolParser(ToolParser): rather than waiting for the complete tag. """ + supports_required_and_named = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) # Stateful streaming fields @@ -132,15 +136,15 @@ class PoolsideV1ToolParser(ToolParser): if tools is None: return False for tool in tools: - if tool.function.name != tool_name: + # ChatCompletion tools nest under .function; Responses + # FunctionTool is flat (.name/.parameters at the top level). + fn = getattr(tool, "function", tool) + if getattr(fn, "name", None) != tool_name: continue - if tool.function.parameters is None: + params = getattr(fn, "parameters", None) + if params is None: return False - arg_type = ( - tool.function.parameters.get("properties", {}) - .get(arg_name, {}) - .get("type", None) - ) + arg_type = params.get("properties", {}).get(arg_name, {}).get("type", None) return arg_type == "string" logger.debug("No tool named '%s'.", tool_name) return False @@ -159,7 +163,19 @@ class PoolsideV1ToolParser(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 super().adjust_request() so it + does not install JSON guided decoding. These models emit XML tool + calls (per the chat template), which JSON guidance would break. + """ + if request.tools: + tc = request.tool_choice + if tc == "required" or isinstance( + tc, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + ): + 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 @@ -192,9 +208,12 @@ class PoolsideV1ToolParser(ToolParser): arg_dct: dict[str, Any] = {} for key, value in pairs: arg_key = key.strip() - arg_val = value.strip() - if not self._is_string_type(tc_name, arg_key, request.tools): - arg_val = self._deserialize(arg_val) + # Keep string values verbatim; whitespace is significant + # (e.g. code/file content). Only strip non-string types. + if self._is_string_type(tc_name, arg_key, request.tools): + arg_val = value + else: + arg_val = self._deserialize(value.strip()) logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val) arg_dct[arg_key] = arg_val tool_calls.append( From 5e3dad04b10df208513d4941da9e72e6d9e77048 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Fri, 26 Jun 2026 15:43:29 +0800 Subject: [PATCH 029/138] [Misc] Move the legacy api_server.py to the examples directory. (#46783) Signed-off-by: wang.yuqi --- docs/design/arch_overview.md | 2 +- docs/examples/README.md | 2 +- .../{chatbot/api_client.py => api_server/client.py} | 4 ++-- .../applications/api_server/server.py | 2 +- examples/applications/chatbot/gradio_webserver.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) rename examples/applications/{chatbot/api_client.py => api_server/client.py} (94%) rename vllm/entrypoints/api_server.py => examples/applications/api_server/server.py (99%) diff --git a/docs/design/arch_overview.md b/docs/design/arch_overview.md index e419104bae3..c19ea49d96d 100644 --- a/docs/design/arch_overview.md +++ b/docs/design/arch_overview.md @@ -178,7 +178,7 @@ incoming requests. The `AsyncLLMEngine` is designed for online serving, where it can handle multiple concurrent requests and stream outputs to clients. The OpenAI-compatible API server uses the `AsyncLLMEngine`. There is also a demo -API server that serves as a simpler example in [vllm/entrypoints/api_server.py](../../vllm/entrypoints/api_server.py). +API server that serves as a simpler example in [examples/applications/api_server/server.py](../../examples/applications/api_server/server.py). The code for `AsyncLLMEngine` can be found in [vllm/engine/async_llm_engine.py](../../vllm/engine/async_llm_engine.py). diff --git a/docs/examples/README.md b/docs/examples/README.md index 9d6126a65c4..5569db9119c 100644 --- a/docs/examples/README.md +++ b/docs/examples/README.md @@ -9,7 +9,7 @@ vLLM's examples are organized into the following categories: - **[`features/`](../../examples/features)** โ€“ Demonstrations of individual vLLM features: automatic prefix caching, speculative decoding, LoRA, structured outputs, prompt embedding, pause/resume, batch invariance, KV events, data parallelism, and more. - **[`reasoning/`](../../examples/reasoning)** โ€“ Examples for reasoning with vLLM. - **[`tool_calling/`](../../examples/tool_calling)** โ€“ Examples for function/tool calling with vLLM. -- **[`applications/`](../../examples/applications)** โ€“ Application examples such as chatbots and RAG (Retrieval-Augmented Generation). +- **[`applications/`](../../examples/applications)** โ€“ Application examples such as simpler api server, chatbots and RAG (Retrieval-Augmented Generation). - **[`rl/`](../../examples/rl)** โ€“ Reinforcement learning examples. - **[`deployment/`](../../examples/deployment)** โ€“ Examples for deploying vLLM in production. - **[`ray_serving/`](../../examples/ray_serving)** โ€“ Scalable serving using Ray. diff --git a/examples/applications/chatbot/api_client.py b/examples/applications/api_server/client.py similarity index 94% rename from examples/applications/chatbot/api_client.py rename to examples/applications/api_server/client.py index 84854911bad..89207d854c9 100644 --- a/examples/applications/chatbot/api_client.py +++ b/examples/applications/api_server/client.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Example Python client for `vllm.entrypoints.api_server` +"""Example Python client for `examples/applications/api_server/server.py` Start the demo server: - python -m vllm.entrypoints.api_server --model + python examples/applications/api_server/server.py --model NOTE: The API server is used only for demonstration and simple performance benchmarks. It is not intended for production use. diff --git a/vllm/entrypoints/api_server.py b/examples/applications/api_server/server.py similarity index 99% rename from vllm/entrypoints/api_server.py rename to examples/applications/api_server/server.py index f950b52d881..adac4133210 100644 --- a/vllm/entrypoints/api_server.py +++ b/examples/applications/api_server/server.py @@ -31,7 +31,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.system_utils import set_ulimit from vllm.version import __version__ as VLLM_VERSION -logger = init_logger("vllm.entrypoints.api_server") +logger = init_logger("api_server") app = FastAPI() engine = None diff --git a/examples/applications/chatbot/gradio_webserver.py b/examples/applications/chatbot/gradio_webserver.py index f75636409c2..005bb7c68c9 100644 --- a/examples/applications/chatbot/gradio_webserver.py +++ b/examples/applications/chatbot/gradio_webserver.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Example for starting a Gradio Webserver Start vLLM API server: - python -m vllm.entrypoints.api_server \ + python examples/applications/api_server/server.py \ --model meta-llama/Llama-2-7b-chat-hf Start Webserver: From bf292b5f6b537d154fc09a3b232f89cbc66827f5 Mon Sep 17 00:00:00 2001 From: AgenticSpark Date: Fri, 26 Jun 2026 01:02:50 -0700 Subject: [PATCH 030/138] [Docs] Remove BambaForCausalLM from supported hybrid models list (#46071) Signed-off-by: liejiang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/usage/v1_guide.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index eca23a11bc8..5613d5ba4e8 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -125,7 +125,7 @@ We are working on enabling prefix caching and chunked prefill for more categorie Models using selective state-space mechanisms instead of standard transformer attention are supported. Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported. -Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`, +Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`). Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`). From d980a3cc6ed9fc83386894211170f1ca85ac9735 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 26 Jun 2026 04:09:56 -0500 Subject: [PATCH 031/138] [ROCm] Fix AITER_UNIFIED_ATTN Dispatching After AITER Bump (#46780) Signed-off-by: Micah Williamson Signed-off-by: Rohan138 Co-authored-by: Rohan138 --- .buildkite/hardware_tests/amd.yaml | 2 ++ docs/design/attention_backends.md | 2 +- tests/compile/passes/test_fusion_attn.py | 5 +++++ .../kernels/attention/test_rocm_aiter_unified_attn.py | 3 ++- .../test_rocm_attention_backends_selection.py | 10 +++++++++- vllm/v1/attention/backends/rocm_aiter_unified_attn.py | 11 +++++++++++ 6 files changed, 30 insertions(+), 3 deletions(-) diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index c2510f38aab..a18241cf18b 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.buildkite/hardware_tests/amd.yaml @@ -6,6 +6,7 @@ steps: # differ ci_base is rebuilt and pushed automatically. - label: "AMD: :docker: ensure ci_base" key: ensure-ci-base-amd + soft_fail: false depends_on: [] device: amd_cpu no_plugin: true @@ -26,6 +27,7 @@ steps: - label: "AMD: :docker: build test image and artifacts" key: image-build-amd + soft_fail: false depends_on: - ensure-ci-base-amd device: amd_cpu diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d268d5b4db2..9278ab6761a 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first). | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | โŒ | โŒ | โŒ | โœ… | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | โŒ | โœ… | โœ… | โŒ | Decoder, Encoder Only | Any | | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | โœ… | โœ… | โŒ | โŒ | Decoder | N/A | -| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | โœ… | โŒ | โœ… | โŒ | All | N/A | +| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | โœ… | โŒ | โœ… | โŒ | All | N/A | | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | โŒ | โœ… | โœ… | โŒ | Decoder, Encoder, Encoder Only | N/A | | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | โœ… | โœ… | โœ… | โŒ | All | Any | | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | โŒ | โŒ | โŒ | โŒ | Decoder | Any | diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index b776f6af98a..531d26e008a 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -306,6 +306,11 @@ def test_attention_quant_pattern( torch.manual_seed(42) backend_cls = backend.get_class() + + # TODO: drop once AITER reenables fp16 unified attention. + if dtype not in backend_cls.supported_dtypes: + pytest.skip(f"{backend.name} does not support dtype {dtype}") + block_size = backend_cls.get_preferred_block_size(16) model_config = ModelConfig( diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index 9e33f24ea28..c02a457c98a 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -30,7 +30,8 @@ NUM_Q_HEADS = 8 NUM_KV_HEADS = 8 HEAD_SIZES = [128, 256] BLOCK_SIZES = [16, 64] -DTYPES = [torch.bfloat16, torch.float16] +# TODO: re-add torch.float16 once AITER reenables fp16 unified attention. +DTYPES = [torch.bfloat16] FP8_DTYPE = current_platform.fp8_dtype() # (query_len, kv_len) per sequence diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 8f9e8acac60..48c6de8f8bd 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -136,9 +136,17 @@ def test_standard_attention_backend_selection( # Get the backend class path from vllm.platforms.rocm import RocmPlatform + # The AITER unified attention kernel only supports BF16/FP8 KV caches + # (its 3D kernel asserts on fp16), so it must be selected with bf16. + dtype = ( + torch.bfloat16 + if selected_backend == "ROCM_AITER_UNIFIED_ATTN" + else torch.float16 + ) + attn_selector_config = AttentionSelectorConfig( head_size=128, - dtype=torch.float16, + dtype=dtype, kv_cache_dtype="auto", block_size=16, use_mla=False, diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 984fc20ecaf..d8363169a8c 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer with PagedAttention and Triton prefix prefill.""" +from typing import ClassVar + import torch from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops +from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -24,6 +27,14 @@ logger = init_logger(__name__) class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend): + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(16)] From 950ee4c2e48fd462a03cebb3283ac9fdcb27b2e4 Mon Sep 17 00:00:00 2001 From: Hyunkyun Moon Date: Fri, 26 Jun 2026 21:02:52 +0900 Subject: [PATCH 032/138] [API] Add token offsets to render endpoints (/v1/.../render) (#44226) Signed-off-by: HyunKyun Moon --- .../openai/test_render_token_offsets.py | 80 +++++++ tests/entrypoints/serve/render/test_render.py | 106 ++++++++++ tests/renderers/test_completions.py | 5 + tests/renderers/test_token_offsets.py | 199 ++++++++++++++++++ .../openai/chat_completion/protocol.py | 16 ++ .../entrypoints/openai/completion/protocol.py | 16 ++ vllm/entrypoints/serve/disagg/protocol.py | 7 + vllm/entrypoints/serve/render/serving.py | 2 + vllm/inputs/engine.py | 4 + vllm/inputs/llm.py | 6 + vllm/renderers/base.py | 89 ++++++-- vllm/renderers/hf.py | 5 + vllm/renderers/params.py | 5 + 13 files changed, 518 insertions(+), 22 deletions(-) create mode 100644 tests/entrypoints/openai/test_render_token_offsets.py create mode 100644 tests/renderers/test_token_offsets.py diff --git a/tests/entrypoints/openai/test_render_token_offsets.py b/tests/entrypoints/openai/test_render_token_offsets.py new file mode 100644 index 00000000000..f7653ab66fc --- /dev/null +++ b/tests/entrypoints/openai/test_render_token_offsets.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the token-offsets request/response protocol wiring: +the request flag flowing into ``TokenizeParams`` and the ``GenerateRequest`` +serialization boundary. End-to-end behavior is covered by +``tests/entrypoints/serve/render/test_render.py``; plain Pydantic field +storage is not retested here. +""" + +from unittest.mock import Mock + +from vllm.config import ModelConfig +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.serve.disagg.protocol import GenerateRequest +from vllm.sampling_params import SamplingParams + + +def _model_config() -> Mock: + model_config = Mock(spec=ModelConfig) + model_config.max_model_len = 128 + return model_config + + +def test_completion_flag_forwarded_to_tok_params(): + """build_tok_params must forward return_token_offsets, defaulting to + False (zero behavioral change for existing callers) and coercing JSON + null to False via the bool() guard.""" + cfg = _model_config() + + default = CompletionRequest(model="m", prompt="hi") + assert default.build_tok_params(cfg).return_token_offsets is False + + on = CompletionRequest(model="m", prompt="hi", return_token_offsets=True) + assert on.build_tok_params(cfg).return_token_offsets is True + + null = CompletionRequest(model="m", prompt="hi", return_token_offsets=None) + assert null.build_tok_params(cfg).return_token_offsets is False + + +def test_chat_flag_forwarded_to_tok_params(): + """Chat build_tok_params has its own (max_completion_tokens) branch, so + its return_token_offsets forwarding is verified independently.""" + cfg = _model_config() + messages = [{"role": "user", "content": "hi"}] + + default = ChatCompletionRequest(model="m", messages=messages) + assert default.build_tok_params(cfg).return_token_offsets is False + + on = ChatCompletionRequest(model="m", messages=messages, return_token_offsets=True) + assert on.build_tok_params(cfg).return_token_offsets is True + + null = ChatCompletionRequest( + model="m", messages=messages, return_token_offsets=None + ) + assert null.build_tok_params(cfg).return_token_offsets is False + + +def test_generate_request_token_offsets_default_none(): + """Defaults to None so existing /v1/.../render responses are unchanged.""" + req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=SamplingParams()) + assert req.token_offsets is None + + +def test_generate_request_token_offsets_survive_json_round_trip(): + """GenerateRequest crosses the disagg serialization boundary; the + tuple[int, int] offsets must survive model_dump and re-validate.""" + req = GenerateRequest( + token_ids=[10, 20], + sampling_params=SamplingParams(), + token_offsets=[(0, 1), (1, 3)], + ) + dumped = req.model_dump() + assert dumped["token_offsets"] == [(0, 1), (1, 3)] + # Re-validate from the dumped dict (sampling_params doesn't round-trip + # cleanly via dump, so re-inject a fresh instance). + again = GenerateRequest.model_validate( + {**dumped, "sampling_params": SamplingParams()} + ) + assert again.token_offsets == [(0, 1), (1, 3)] diff --git a/tests/entrypoints/serve/render/test_render.py b/tests/entrypoints/serve/render/test_render.py index 7aacf4564e3..d7339361ff7 100644 --- a/tests/entrypoints/serve/render/test_render.py +++ b/tests/entrypoints/serve/render/test_render.py @@ -263,3 +263,109 @@ async def test_chat_completion_render_with_sampling_params(client): # Check that internal fields are not present assert "_all_stop_token_ids" not in sampling_params + + +@pytest.mark.asyncio +async def test_completion_render_emits_token_offsets(client): + """With return_token_offsets, /v1/completions/render returns per-token + (start, end) char offsets aligned with token_ids.""" + prompt = "Hello, world." + response = await client.post( + "/v1/completions/render", + json={ + "model": MODEL_NAME, + "prompt": prompt, + "return_token_offsets": True, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert isinstance(data, list) + offsets = data[0]["token_offsets"] + assert offsets is not None + assert len(offsets) == len(data[0]["token_ids"]) + for start, end in offsets: + assert isinstance(start, int) and isinstance(end, int) + assert 0 <= start <= end <= len(prompt) + + +@pytest.mark.asyncio +async def test_completion_render_default_no_token_offsets(client): + """Without the flag, token_offsets must be null (existing responses + unchanged).""" + response = await client.post( + "/v1/completions/render", + json={ + "model": MODEL_NAME, + "prompt": "Hello, world.", + }, + ) + + assert response.status_code == 200 + data = response.json() + assert data[0]["token_offsets"] is None + + +@pytest.mark.asyncio +async def test_chat_render_emits_token_offsets(client): + """With return_token_offsets, /v1/chat/completions/render returns + per-token offsets relative to the templated prompt string.""" + response = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [{"role": "user", "content": "Hello, world."}], + "return_token_offsets": True, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert isinstance(data, dict) + offsets = data["token_offsets"] + assert offsets is not None + assert len(offsets) == len(data["token_ids"]) + for start, end in offsets: + assert isinstance(start, int) and isinstance(end, int) + assert 0 <= start <= end + + +@pytest.mark.asyncio +async def test_chat_render_default_no_token_offsets(client): + """Without the flag, chat render token_offsets must be null.""" + response = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [{"role": "user", "content": "Hello, world."}], + }, + ) + + assert response.status_code == 200 + data = response.json() + assert data["token_offsets"] is None + + +@pytest.mark.asyncio +async def test_completion_render_multiple_prompts_token_offsets(client): + """Each prompt in a batch gets its own offsets aligned with its tokens.""" + prompts = ["Hello, world.", "Goodbye, world."] + response = await client.post( + "/v1/completions/render", + json={ + "model": MODEL_NAME, + "prompt": prompts, + "return_token_offsets": True, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data) == len(prompts) + for item, prompt in zip(data, prompts): + offsets = item["token_offsets"] + assert offsets is not None + assert len(offsets) == len(item["token_ids"]) + for start, end in offsets: + assert 0 <= start <= end <= len(prompt) diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 00d604afdcf..76e88f4213e 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -79,6 +79,11 @@ class DummyTokenizer: return list(range(in_length)) + def __call__(self, text: str, **kwargs): + # BaseRenderer._tokenize_prompt calls the tokenizer via __call__ (to + # unify the output type), so mirror a real tokenizer's BatchEncoding. + return {"input_ids": self.encode(text, **kwargs)} + def _build_renderer( model_config: MockModelConfig, diff --git a/tests/renderers/test_token_offsets.py b/tests/renderers/test_token_offsets.py new file mode 100644 index 00000000000..ab881782659 --- /dev/null +++ b/tests/renderers/test_token_offsets.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for renderer-level token-offset behavior. + +These exercise ``_tokenize_prompt`` (offset extraction + capability/MM +gating) and the ``_tokenize_prompt -> _process_tokens -> TokensInput`` +forwarding chain. Endpoint-level coverage lives in +``tests/entrypoints/serve/render/test_render.py``. +""" + +import pytest + +from vllm.renderers.params import TokenizeParams + + +@pytest.fixture +def fast_tokenizer(): + """gpt2 ships a Fast tokenizer; use it to test the offsets happy path.""" + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained("openai-community/gpt2", use_fast=True) + + +def _make_base_renderer_with(tokenizer): + """Build a minimal BaseRenderer subclass that exposes the tokenizer so we + can call ``_tokenize_prompt`` directly. BaseRenderer is abstract because of + ``render_messages``; we just need a stub.""" + from vllm.renderers.base import BaseRenderer + + class _StubRenderer(BaseRenderer): + def __init__(self, tok): + # Bypass BaseRenderer.__init__ โ€” we don't need a VllmConfig. + from vllm.utils.async_utils import make_async + + self.tokenizer = tok + self._executor = None + # Mirror BaseRenderer.__init__: the async path offloads the sync + # ``_tokenize_prompt`` to a thread pool. + self._tokenize_prompt_async = make_async(self._tokenize_prompt) + self.mm_processor = None + + def get_tokenizer(self): + return self.tokenizer + + def _can_produce_offsets(self): + # Mirror HfRenderer: offsets only for fast tokenizers. + return self.tokenizer is not None and self.tokenizer.is_fast + + def render_messages(self, messages, params): # pragma: no cover + raise NotImplementedError + + return _StubRenderer(tokenizer) + + +class TestTokenizePromptOffsets: + def test_fast_tokenizer_with_flag_returns_offsets(self, fast_tokenizer): + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + prompt = {"prompt": "Hello, world."} + + result = renderer._tokenize_prompt(prompt, params) + + assert "prompt_token_ids" in result + offsets = result["prompt_token_offsets"] + assert offsets is not None + # Length must match the token sequence, and each (start, end) is an + # ordered pair within the source text. + assert len(offsets) == len(result["prompt_token_ids"]) + text_len = len("Hello, world.") + for s, e in offsets: + assert isinstance(s, int) and isinstance(e, int) + assert 0 <= s <= e <= text_len + + def test_base_renderer_without_override_yields_no_offsets(self, fast_tokenizer): + """A renderer that does not override ``_can_produce_offsets`` never + emits offsets, even with a fast tokenizer and the flag set. This locks + in the base-default-False / subclass-override design.""" + from vllm.renderers.base import BaseRenderer + + class _BareRenderer(BaseRenderer): + def __init__(self, tok): + self.tokenizer = tok + self._executor = None + self.mm_processor = None + + def get_tokenizer(self): + return self.tokenizer + + def render_messages(self, messages, params): # pragma: no cover + raise NotImplementedError + + renderer = _BareRenderer(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + + result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) + + assert "prompt_token_offsets" not in result + + def test_default_flag_no_offsets(self, fast_tokenizer): + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None) # flag defaults False + + result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) + + # Field must be absent (not None) so TokensInput serialization stays + # minimal for existing consumers. + assert "prompt_token_offsets" not in result + + def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer): + """Force is_fast=False to simulate a Slow tokenizer: the flag is set + but offsets must not be returned because it cannot produce them.""" + from unittest.mock import PropertyMock, patch + + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + + with patch.object( + type(fast_tokenizer), + "is_fast", + new_callable=PropertyMock, + return_value=False, + ): + result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) + + assert "prompt_token_offsets" not in result + + @pytest.mark.parametrize("mm_key", ["multi_modal_data", "multi_modal_uuids"]) + def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key): + """Offsets index the text prompt, which is meaningless once multimodal + data is interleaved, so they are suppressed when MM inputs are present.""" + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + prompt = {"prompt": "Hello.", mm_key: {"image": ["x"]}} + + result = renderer._tokenize_prompt(prompt, params) + + assert "prompt_token_offsets" not in result + + @pytest.mark.asyncio + async def test_tokenize_prompt_async_returns_offsets(self, fast_tokenizer): + """The async path offloads the sync tokenizer; it must yield the same + offsets as the sync path.""" + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + + result = await renderer._tokenize_prompt_async( + {"prompt": "Hello, world."}, params + ) + + offsets = result["prompt_token_offsets"] + assert offsets is not None + assert len(offsets) == len(result["prompt_token_ids"]) + + +class TestProcessTokensForwardsOffsets: + """Tests that the ``_tokenize_prompt -> _process_tokens -> TokensInput`` + chain carries ``prompt_token_offsets`` through to the engine input. + ``_process_tokens`` rebuilds the engine input from scratch, so it must + copy the field explicitly. The sync and async variants are independent + implementations, so both are checked. + """ + + def test_sync_forwards_offsets_to_engine_input(self, fast_tokenizer): + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + + tokens_prompt = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) + # Sanity: offsets must reach the TokensPrompt, else this guards the + # wrong layer. + expected = tokens_prompt["prompt_token_offsets"] + + engine_input = renderer._process_tokens(tokens_prompt) + + assert engine_input["prompt_token_offsets"] == expected + + @pytest.mark.asyncio + async def test_async_forwards_offsets_to_engine_input(self, fast_tokenizer): + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) + + tokens_prompt = await renderer._tokenize_prompt_async( + {"prompt": "Hello, world."}, params + ) + expected = tokens_prompt["prompt_token_offsets"] + + engine_input = await renderer._process_tokens_async(tokens_prompt) + + assert engine_input["prompt_token_offsets"] == expected + + def test_no_offsets_forwarded_when_flag_off(self, fast_tokenizer): + renderer = _make_base_renderer_with(fast_tokenizer) + params = TokenizeParams(max_total_tokens=None) # flag defaults False + + tokens_prompt = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) + assert "prompt_token_offsets" not in tokens_prompt + + engine_input = renderer._process_tokens(tokens_prompt) + + assert "prompt_token_offsets" not in engine_input diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 09ce8bf8dab..aa2af69777c 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -382,6 +382,21 @@ class ChatCompletionRequest(OpenAIBaseModel): "need to map generated text back to input tokens." ), ) + return_token_offsets: bool | None = Field( + default=False, + description=( + "If true, return char-level (start, end) offsets for each " + "token relative to the tokenized source string in the " + "`token_offsets` field of the rendered response. Only " + "supported on the `/v1/completions/render` and " + "`/v1/chat/completions/render` endpoints; ignored on regular " + "generation endpoints. Honored only for Fast (Rust-backed) " + "tokenizers; otherwise `token_offsets` is null. For chat " + "requests, offsets are relative to the templated prompt " + "string (after applying the chat template). Multimodal " + "inputs and pre-tokenized inputs always yield null." + ), + ) return_prompt_text: bool | None = Field( default=None, description=( @@ -524,6 +539,7 @@ class ChatCompletionRequest(OpenAIBaseModel): needs_detokenization=bool(self.echo and not self.return_token_ids), max_total_tokens_param="max_model_len", max_output_tokens_param=max_output_tokens_param, + return_token_offsets=bool(self.return_token_offsets), ) # Default sampling parameters for chat completion requests diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 1d61ca3c598..b5b715b50bd 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -152,6 +152,21 @@ class CompletionRequest(OpenAIBaseModel): "need to map generated text back to input tokens." ), ) + return_token_offsets: bool | None = Field( + default=False, + description=( + "If true, return char-level (start, end) offsets for each " + "token relative to the tokenized source string in the " + "`token_offsets` field of the rendered response. Only " + "supported on the `/v1/completions/render` and " + "`/v1/chat/completions/render` endpoints; ignored on regular " + "generation endpoints. Honored only for Fast (Rust-backed) " + "tokenizers; otherwise `token_offsets` is null. For chat " + "requests, offsets are relative to the templated prompt " + "string (after applying the chat template). Multimodal " + "inputs and pre-tokenized inputs always yield null." + ), + ) cache_salt: str | None = Field( default=None, @@ -209,6 +224,7 @@ class CompletionRequest(OpenAIBaseModel): needs_detokenization=bool(self.echo and not self.return_token_ids), max_total_tokens_param="max_model_len", max_output_tokens_param="max_tokens", + return_token_offsets=bool(self.return_token_offsets), ) # Default sampling parameters for completion requests diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index 2e98f5e811c..d20752a9063 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -82,6 +82,13 @@ class GenerateRequest(BaseModel): raise ValueError("token_ids must not contain negative values") return v + token_offsets: list[tuple[int, int]] | None = None + """Char-level (start, end) offsets per token, relative to the + tokenized source string. Present only when the request set + `return_token_offsets=True` and the renderer was able to compute + them (Fast tokenizer, text input, no multimodal data). List length + equals `token_ids` length when present. None otherwise.""" + features: MultiModalFeatures | None = None """Multimodal hashes and placeholder positions (populated for MM inputs).""" diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 42cf2460c41..1bba26722b9 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -139,6 +139,7 @@ class ServingRender(BaseServing): stream_options=(request.stream_options if request.stream else None), cache_salt=request.cache_salt, priority=request.priority, + token_offsets=engine_input.get("prompt_token_offsets"), ) async def render_completion_request( @@ -194,6 +195,7 @@ class ServingRender(BaseServing): stream_options=(request.stream_options if request.stream else None), cache_salt=request.cache_salt, priority=request.priority, + token_offsets=engine_input.get("prompt_token_offsets"), ) ) diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index 1c12fbc2c55..eacadcbc924 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -38,6 +38,10 @@ class TokensInput(_InputOptions): prompt: NotRequired[str] """The prompt text corresponding to the token IDs, if available.""" + prompt_token_offsets: NotRequired[list[tuple[int, int]] | None] + """Char-level (start, end) offsets per token, propagated from the + renderer's TokensPrompt when offsets were computed.""" + def tokens_input( prompt_token_ids: list[int], diff --git a/vllm/inputs/llm.py b/vllm/inputs/llm.py index 918098b758c..f03661078c1 100644 --- a/vllm/inputs/llm.py +++ b/vllm/inputs/llm.py @@ -115,6 +115,12 @@ class TokensPrompt(_PromptOptions): token_type_ids: NotRequired[list[int]] """A list of token type IDs to pass to the cross encoder model.""" + prompt_token_offsets: NotRequired[list[tuple[int, int]] | None] + """Char-level (start, end) offsets per token, relative to the + tokenized source string. Present only when offsets were requested + AND a Fast (Rust-backed) tokenizer was used AND no multimodal data + was present. The list length equals the length of `prompt_token_ids`.""" + class EmbedsPrompt(_PromptOptions): """Schema for a prompt provided via token embeddings.""" diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 9f4794faa0d..00cbec33d6f 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -89,8 +89,13 @@ class BaseRenderer(ABC, Generic[_T]): # to keep the asyncio event loop responsive under concurrent load. self._mm_executor: Executor = self._executor - # Offloading tokenizer encode & decode to thread pool. - self._async_tokenizer_encode = make_async(self._encode, executor=self._executor) + # Offload tokenization to the thread pool. The sync + # ``_tokenize_prompt`` already encapsulates the unified ``__call__`` + # path and char-offset extraction, so the async variant is just it + # offloaded (mirrors ``_process_multimodal_async`` below). + self._tokenize_prompt_async = make_async( + self._tokenize_prompt, executor=self._executor + ) self._async_tokenizer_decode = make_async(self._decode, executor=self._executor) self.mm_processor: BaseMultiModalProcessor | None = None @@ -147,9 +152,6 @@ class BaseRenderer(ABC, Generic[_T]): def _decode(self, *args, **kwargs): return self.get_tokenizer().decode(*args, **kwargs) - def _encode(self, *args, **kwargs): - return self.get_tokenizer().encode(*args, **kwargs) - def get_mm_processor(self) -> "BaseMultiModalProcessor": if self.mm_processor is None: raise ValueError("Multi-modal processor not available for text-only models") @@ -414,31 +416,64 @@ class BaseRenderer(ABC, Generic[_T]): return self.render_messages(messages, params) # Step 2: Tokenize prompts if necessary + def _can_produce_offsets(self) -> bool: + """Whether this renderer's tokenizer can emit char-level offsets. + + Defaults to False; only renderers backed by an HF fast tokenizer + (see ``HfRenderer``) can produce ``offset_mapping``. + """ + return False + + def _wants_offsets( + self, + prompt: "TextPrompt", + params: "TokenizeParams", + ) -> bool: + return ( + params.return_token_offsets + and self._can_produce_offsets() + and not prompt.get("multi_modal_data") + and not prompt.get("multi_modal_uuids") + ) + + @staticmethod + def _build_tokens_prompt( + token_ids: Sequence[int], + prompt: "TextPrompt", + *, + offset_mapping: Sequence[tuple[int, int]] | None = None, + ) -> "TokensPrompt": + """Build a TokensPrompt from already-extracted token ids. + + ``offset_mapping`` is the per-token ``(start, end)`` sequence from + a BatchEncoding; pass it only when offsets were requested, and it + is attached as ``prompt_token_offsets``. + """ + if offset_mapping is not None: + return TokensPrompt( + prompt_token_ids=list(token_ids), + prompt_token_offsets=[(int(s), int(e)) for s, e in offset_mapping], + **prompt, + ) + return TokensPrompt(prompt_token_ids=list(token_ids), **prompt) + def _tokenize_prompt( self, prompt: TextPrompt, params: TokenizeParams, ) -> TokensPrompt: tokenizer = self.get_tokenizer() - prompt_token_ids = tokenizer.encode( - prompt["prompt"], - **params.get_encode_kwargs(), + want_offsets = self._wants_offsets(prompt, params) + kwargs = params.get_encode_kwargs() + if want_offsets: + kwargs = {**kwargs, "return_offsets_mapping": True} + encoding = tokenizer(prompt["prompt"], **kwargs) + return self._build_tokens_prompt( + encoding["input_ids"], + prompt, + offset_mapping=encoding["offset_mapping"] if want_offsets else None, ) - return TokensPrompt(prompt_token_ids=prompt_token_ids, **prompt) - - async def _tokenize_prompt_async( - self, - prompt: TextPrompt, - params: TokenizeParams, - ) -> TokensPrompt: - prompt_token_ids = await self._async_tokenizer_encode( - prompt["prompt"], - **params.get_encode_kwargs(), - ) - - return TokensPrompt(prompt_token_ids=prompt_token_ids, **prompt) - def _detokenize_prompt(self, prompt: TokensPrompt) -> TokensPrompt: tokenizer = self.get_tokenizer() prompt["prompt"] = tokenizer.decode(prompt["prompt_token_ids"]) @@ -747,6 +782,11 @@ class BaseRenderer(ABC, Generic[_T]): engine_input["prompt"] = prompt_text if cache_salt := prompt.get("cache_salt"): engine_input["cache_salt"] = cache_salt + # Narrow the union โ€” `prompt_token_offsets` is only on TokensInput. + if engine_input["type"] == "token" and ( + (offsets := prompt.get("prompt_token_offsets")) is not None + ): + engine_input["prompt_token_offsets"] = offsets return engine_input @@ -805,6 +845,11 @@ class BaseRenderer(ABC, Generic[_T]): engine_input["prompt"] = prompt_text if cache_salt := prompt.get("cache_salt"): engine_input["cache_salt"] = cache_salt + # Narrow the union โ€” `prompt_token_offsets` is only on TokensInput. + if engine_input["type"] == "token" and ( + (offsets := prompt.get("prompt_token_offsets")) is not None + ): + engine_input["prompt_token_offsets"] = offsets return engine_input diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index e57d0586aa0..ea0902c8806 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -882,6 +882,11 @@ class HfRenderer(BaseRenderer[HfTokenizer]): self.tokenizer, config.model_config.renderer_num_workers + 1 ) + def _can_produce_offsets(self) -> bool: + # HF tokenizers may be slow (use_fast=False); only fast tokenizers + # expose offset_mapping. + return self.tokenizer is not None and self.tokenizer.is_fast + def render_messages( self, messages: list[ChatCompletionMessageParam], diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index d5c89abc043..8e0aaf303cc 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -167,6 +167,11 @@ class TokenizeParams: add_special_tokens: bool = True """Whether to add special tokens.""" + return_token_offsets: bool = False + """If true, request char-level (start, end) offsets per token. Honored + only for Fast (Rust-backed) tokenizers with text input and no multimodal + data; otherwise silently ignored.""" + needs_detokenization: bool = False """ Whether the tokenized prompt needs to contain the original text. From 302954e5f603b30a8fe6d4c84b7e655f0e3e74db Mon Sep 17 00:00:00 2001 From: TJian Date: Fri, 26 Jun 2026 21:33:35 +0800 Subject: [PATCH 033/138] [ROCm] [CI] fix transcription flakiness AMD: Entrypoints Integration (API Server OpenAI - Part 1) (mi325_1) (#46823) Signed-off-by: tjtanaa --- tests/entrypoints/openai/test_run_batch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/entrypoints/openai/test_run_batch.py b/tests/entrypoints/openai/test_run_batch.py index cd1daf0bbbc..0f7d7f5f464 100644 --- a/tests/entrypoints/openai/test_run_batch.py +++ b/tests/entrypoints/openai/test_run_batch.py @@ -305,6 +305,7 @@ INPUT_TRANSCRIPTION_HTTP_BATCH = ( "body": { "model": SPEECH_LARGE_MODEL_NAME, "file_url": AudioAsset("mary_had_lamb").url, + "language": "en", "response_format": "json", }, } From 8e394244a59afc67a37bf47dab0ab76bf5ce5885 Mon Sep 17 00:00:00 2001 From: qli88 Date: Fri, 26 Jun 2026 08:35:35 -0500 Subject: [PATCH 034/138] [ROCm]Enable AITER MoE backend for MiniMax-M3-MXFP4 (#46419) Signed-off-by: Qiang Li Co-authored-by: TJian --- .../model_executor/layers/fused_moe/config.py | 2 ++ .../fused_moe/experts/rocm_aiter_moe.py | 30 ++++++++++++++----- vllm/model_executor/layers/fused_moe/layer.py | 2 ++ vllm/models/minimax_m3/amd/model.py | 1 + 4 files changed, 27 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index d7ad59b46cd..b065d2142ce 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1279,6 +1279,8 @@ class FusedMoEConfig: hidden_dim_unpadded: int | None = None # Defaults to intermediate_size_per_partition if not specified. intermediate_size_per_partition_unpadded: int | None = None + # Model specific override + intermediate_pad: int | None = None moe_backend: MoEBackend = "auto" max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index bd9b285fe74..4f191334bb2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -251,12 +251,17 @@ def rocm_aiter_fused_experts( if quant_config is None: quant_config = FUSED_MOE_UNQUANTIZED_CONFIG + # Gate/up interleave hint; only the SWIGLUOAI activations override it. + activation_interleave = None if activation == MoEActivation.SILU: activation_method = ActivationMethod.SILU elif activation == MoEActivation.GELU: activation_method = ActivationMethod.GELU elif activation == MoEActivation.SWIGLUOAI: activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu") + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu") + activation_interleave = False else: raise ValueError(f"Unsupported activation: {activation}") @@ -337,9 +342,14 @@ def rocm_aiter_fused_experts( assert moe_config.intermediate_size_per_partition_unpadded is not None hidden_pad = hidden_states.shape[1] - moe_config.hidden_dim_unpadded intermediate_pad = ( - moe_config.intermediate_size_per_partition - - moe_config.intermediate_size_per_partition_unpadded + ( + moe_config.intermediate_size_per_partition + - moe_config.intermediate_size_per_partition_unpadded + ) + if moe_config.intermediate_pad is None + else moe_config.intermediate_pad ) + # Round hidden_pad/intermediate_pad to match AITER's CK/FlyDSL MoE # dispatch (currently pinned to v0.1.13.post1): # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1073 @@ -357,14 +367,17 @@ def rocm_aiter_fused_experts( # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`, # which always sets `is_guinterleave=True`. # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling. + from aiter.ops.flydsl.moe_common import GateMode + gate_mode = "" if quant_config.use_mxfp4_w4a16: - try: - from aiter.ops.flydsl.moe_common import GateMode - - gate_mode = GateMode.INTERLEAVE.value - except ImportError: - pass + gate_mode = GateMode.INTERLEAVE.value + elif activation_interleave is not None: + gate_mode = ( + GateMode.INTERLEAVE.value + if activation_interleave + else GateMode.SEPARATED.value + ) return rocm_aiter_ops.fused_moe( hidden_states, @@ -458,6 +471,7 @@ class AiterExperts(mk.FusedMoEExpertsModular): MoEActivation.SILU, MoEActivation.GELU, MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, ] @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 22548438586..e92cb7cad09 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -105,6 +105,7 @@ def FusedMoE( top_k: int, hidden_size: int, intermediate_size: int, + intermediate_pad: int | None = None, params_dtype: torch.dtype | None = None, renormalize: bool = True, use_grouped_topk: bool = False, @@ -311,6 +312,7 @@ def FusedMoE( experts_per_token=top_k, hidden_dim=hidden_size, intermediate_size=intermediate_size, + intermediate_pad=intermediate_pad, num_local_experts=expert_map_manager.local_num_experts, num_logical_experts=logical_num_experts, moe_parallel_config=moe_parallel_config, diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 4f3528806c9..f01ed45dd65 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -309,6 +309,7 @@ class MiniMaxM3MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, + intermediate_pad=0, scoring_func=config.scoring_func, e_score_correction_bias=self.e_score_correction_bias, renormalize=True, From 8921c4be88effbd295dd7c2410fd21411256f819 Mon Sep 17 00:00:00 2001 From: TJian Date: Fri, 26 Jun 2026 21:43:27 +0800 Subject: [PATCH 035/138] [ROCm] [Performance] Optimize aiter moe for DeepSeekV4 (#46122) Signed-off-by: tjtanaa --- .../layers/fused_moe/oracle/mxfp4.py | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index e6e9d17925a..56a7a6482d1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1429,38 +1429,51 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: - from vllm._aiter_ops import rocm_aiter_ops # noqa: F401 + # Initially introduced for DeepSeekV4 if w13_bias is not None: w13_bias = w13_bias.data.to(torch.float32) if w2_bias is not None: w2_bias = w2_bias.data.to(torch.float32) - e, n, k = w13_weight.shape + import os - # No de-interleave: standard _load_w13 already produces - # [gate_all, up_all] layout. Use aiter-native shuffle functions - # (matching aiter/ops/flydsl/test_flydsl_moe_a4w4.py pattern). + from aiter.ops.shuffle import shuffle_scale as _shuf_s from aiter.ops.shuffle import shuffle_weight as _shuf_w - from aiter.utility.fp4_utils import e8m0_shuffle as _e8m0_shuf - # w13 (gate+up, stage1): shuffle_weight with layout (16,16) + # TODO: Remove this once AITER is fixed + # Necessary for AITER side from crashing + os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" + w13_weight = torch.nn.Parameter( - _shuf_w(w13_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)), + _shuf_w( + w13_weight.data.view(torch.float4_e2m1fn_x2), + is_guinterleave=True, + gate_up=True, + ), requires_grad=False, ) - shuffled_w13_scale = _e8m0_shuf( - w13_weight_scale.view(-1, w13_weight_scale.shape[-1]) + shuffled_w13_scale = _shuf_s( + w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]), + num_experts, + True, + True, ) - # w2 (down-proj, stage2): same shuffle as w13 for a4w4 fp4x2 - # (tuning script uses shuffle_weight((16,16)) + e8m0_shuffle for both) w2_weight = torch.nn.Parameter( - _shuf_w(w2_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)), + _shuf_w( + w2_weight.data.view(torch.float4_e2m1fn_x2), + is_guinterleave=True, + gate_up=False, + ), requires_grad=False, ) - shuffled_w2_scale = _e8m0_shuf( - w2_weight_scale.view(-1, w2_weight_scale.shape[-1]) + # use_gu_interleave + shuffled_w2_scale = _shuf_s( + w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]), + num_experts, + True, + False, ) return ( From c2507fb2937aa8c8e74bea15719d04fb6090befe Mon Sep 17 00:00:00 2001 From: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Date: Fri, 26 Jun 2026 10:05:20 -0400 Subject: [PATCH 036/138] [ROCm] [MoE] [Perf] Shared-expert fusion for bias-routed MoE; enable on MiniMax-M3 mxfp8 model (#46545) Signed-off-by: Hongxia Yang Co-authored-by: Claude Opus 4.8 --- .../fused_moe/experts/mxfp8_native_moe.py | 8 ++- vllm/model_executor/layers/fused_moe/layer.py | 43 ++++++++++------ .../router/fused_topk_bias_router.py | 26 ++++++++++ .../layers/fused_moe/router/router_factory.py | 3 ++ vllm/models/minimax_m3/amd/model.py | 50 +++++++++++++++++-- 5 files changed, 110 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index fa6e902396f..b511e368f4a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -227,10 +227,16 @@ def fused_moe_mxfp8_native( tiles = _mxfp8_moe_tiles(T) block_m = tiles["block_m"] + # Bin by the actual number of expert weight rows. With fused shared experts + # the weight tensor has more rows than ``global_num_experts`` (the routed + # count), and their ids fall outside [0, global_num_experts); binning by the + # routed count would treat them as invalid. Under EP (expert_map set) the + # tensor holds only local experts, so keep the global count for remapping. + num_align_experts = w13.shape[0] if expert_map is None else global_num_experts sorted_ids, expert_ids, num_post = moe_align_block_size( topk_ids, block_m, - global_num_experts, + num_align_experts, expert_map, ignore_invalid_experts=expert_map is not None, ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index e92cb7cad09..871f905badc 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -6,6 +6,7 @@ from typing import Any import torch +import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import ParallelConfig, get_current_vllm_config from vllm.distributed import ( @@ -77,24 +78,20 @@ def determine_expert_counts( ) -> tuple[int, int, int]: global_num_experts = num_experts + num_redundant_experts logical_num_experts = num_experts - # ROCm aiter shared experts fusion - # AITER only supports gated activations (silu/gelu), so disable it - # for non-gated MoE (is_act_and_mul=False) - # rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul - aiter_fmoe_shared_expert_enabled = ( - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and is_act_and_mul - ) + # Shared-expert fusion: append the shared expert(s) as routed-expert slots + # so they run in the same grouped GEMM. Gated by + # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: either the native aiter fused-MoE + # path (env + master switch, via is_fusion_moe_shared_experts_enabled) or the + # backend-neutral router-append path (env alone, independent of the master + # switch; e.g. the MM3 triton/flydsl mxfp8 MoE). Gated activations only. + fuse_shared_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + or envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS + ) and is_act_and_mul num_fused_shared_experts = ( - n_shared_experts - if n_shared_experts is not None and aiter_fmoe_shared_expert_enabled - else 0 + n_shared_experts if n_shared_experts is not None and fuse_shared_enabled else 0 ) - if not aiter_fmoe_shared_expert_enabled and num_fused_shared_experts != 0: - raise ValueError( - "n_shared_experts is only supported on ROCm aiter when " - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" - ) return global_num_experts, logical_num_experts, num_fused_shared_experts @@ -188,7 +185,8 @@ def FusedMoE( has_bias: Whether expert layers have bias terms is_sequence_parallel: Whether sequence parallelism is enabled expert_mapping: Expert parameter mapping for weight loading - n_shared_experts: Number of shared experts (ROCm aiter only) + n_shared_experts: Number of shared experts to fuse into the routed + grouped GEMM (ROCm; requires aiter FSE or the router-append path) router_logits_dtype: Data type for router logits buffers gate: Pre-configured gate module shared_experts: Pre-configured shared experts module @@ -289,6 +287,19 @@ def FusedMoE( else 1.0, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=num_fused_shared_experts, + # Fused shared-expert slot weight. With apply_routed_scale_to_output + # the runner scales the combined output by routed_scaling_factor, so + # the shared slot weight must be 1/routed_scaling_factor for its net + # contribution to be 1.0 (matching the un-scaled separate-MLP add). + shared_expert_weight=( + (1.0 / routed_scaling_factor) + if ( + apply_routed_scale_to_output + and num_fused_shared_experts > 0 + and routed_scaling_factor + ) + else 1.0 + ), zero_expert_type=zero_expert_type, num_logical_experts=logical_num_experts, hash_indices_table=hash_indices_table, diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index f30f81a53c7..d505c5ce4b7 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -334,6 +334,8 @@ class FusedTopKBiasRouter(BaseRouter): *, scoring_func: str = "sigmoid", hash_indices_table: torch.Tensor | None = None, + num_fused_shared_experts: int = 0, + shared_expert_weight: float = 1.0, ): super().__init__( top_k=top_k, @@ -346,6 +348,11 @@ class FusedTopKBiasRouter(BaseRouter): self.routed_scaling_factor = routed_scaling_factor self.scoring_func = scoring_func self._hash_indices_table = hash_indices_table + # Fused shared experts: append constant slots (ids immediately after + # the routed experts, [global, global+n)) routed to by every token at + # ``shared_expert_weight``, AFTER the routed top-k is renormalized. + self.num_fused_shared_experts = num_fused_shared_experts + self.shared_expert_weight = shared_expert_weight @property def routing_method_type(self) -> RoutingMethodType: @@ -382,4 +389,23 @@ class FusedTopKBiasRouter(BaseRouter): routed_scaling_factor=self.routed_scaling_factor, ) + if self.num_fused_shared_experts > 0: + m = topk_ids.shape[0] + n = self.num_fused_shared_experts + # global_num_experts counts only the routed experts; the fused + # shared experts occupy the slots immediately after them, i.e. ids + # [global_num_experts, global_num_experts + n). + base = self.global_num_experts + shared_ids = torch.arange( + base, base + n, dtype=topk_ids.dtype, device=topk_ids.device + ).expand(m, n) + shared_w = torch.full( + (m, n), + self.shared_expert_weight, + dtype=topk_weights.dtype, + device=topk_weights.device, + ) + topk_ids = torch.cat([topk_ids, shared_ids], dim=-1) + topk_weights = torch.cat([topk_weights, shared_w], dim=-1) + return topk_weights, topk_ids diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index 7246185f394..c7cfccbe64b 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -47,6 +47,7 @@ def create_fused_moe_router( topk_group: int | None = None, scoring_func: str = "softmax", num_fused_shared_experts: int = 0, + shared_expert_weight: float = 1.0, # grouped topk + fused topk bias parameters routed_scaling_factor: float = 1.0, e_score_correction_bias: torch.Tensor | None = None, @@ -188,6 +189,8 @@ def create_fused_moe_router( routed_scaling_factor=routed_scaling_factor, scoring_func=scoring_func, hash_indices_table=hash_indices_table, + num_fused_shared_experts=num_fused_shared_experts, + shared_expert_weight=shared_expert_weight, ) if ( diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index f01ed45dd65..894550c1576 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -23,6 +23,7 @@ import torch from torch import nn from transformers import PretrainedConfig +import vllm.envs as envs from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import ( @@ -95,6 +96,7 @@ from vllm.models.minimax_m3.common.sparse_attention import ( ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype from vllm.v1.kv_cache_interface import ( @@ -104,6 +106,23 @@ from vllm.v1.kv_cache_interface import ( ) +def _fuse_shared_experts_enabled(config: PretrainedConfig) -> bool: + """Whether to fuse the shared expert into the routed grouped MoE. + + ROCm only. Opt-in via ``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS`` (the + router-append fusion runs on the triton/flydsl mxfp8 MoE independent of the + aiter master switch); requires a shared expert and is disabled under expert + parallelism (the shared slot is appended to the routed top-k, which the EP + expert-map path does not handle). + """ + return bool( + current_platform.is_rocm() + and getattr(config, "n_shared_experts", None) + and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS + and not get_current_vllm_config().parallel_config.enable_expert_parallel + ) + + def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]: """Layer ids whose attention runs the extra sparse "index" branch.""" cfg = getattr(config, "sparse_attention_config", None) @@ -294,8 +313,13 @@ class MiniMaxM3MoE(nn.Module): prefix=f"{prefix}.gate", ) + # Fuse the shared expert into the routed grouped GEMM when opted in via + # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: it becomes routed-expert slot + # ``num_local_experts``, reached by every token, eliminating the + # separate dense-MLP launches. Not supported under expert parallelism. + self.fuse_shared_experts = _fuse_shared_experts_enabled(config) self.shared_experts: MiniMaxM3MLP | None = None - if self.n_shared_experts: + if self.n_shared_experts and not self.fuse_shared_experts: self.shared_experts = MiniMaxM3MLP( config=config, intermediate_size=config.intermediate_size * self.n_shared_experts, @@ -321,6 +345,9 @@ class MiniMaxM3MoE(nn.Module): apply_routed_scale_to_output=True, router_logits_dtype=self.gate.out_dtype, shared_experts=self.shared_experts, + n_shared_experts=( + self.n_shared_experts if self.fuse_shared_experts else None + ), quant_config=quant_config, prefix=f"{prefix}.experts", ) @@ -864,13 +891,18 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): return hidden_states def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - # Checkpoint experts use w1=gate, w2=down, w3=up. + # Checkpoint experts use w1=gate, w2=down, w3=up. When fusing the shared + # expert, include the appended slot (id == num_local_experts). + n_shared = getattr(self.config, "n_shared_experts", 0) or 0 + num_experts = self.config.num_local_experts + ( + n_shared if _fuse_shared_experts_enabled(self.config) else 0 + ) return fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", - num_experts=self.config.num_local_experts, + num_experts=num_experts, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -897,6 +929,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + _fuse_shared = _fuse_shared_experts_enabled(self.config) for name, loaded_weight in weights: # The MTP module is not modeled yet. if "mtp." in name: @@ -907,6 +940,17 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): if "weight_scale_inv" in name: name = name.replace("weight_scale_inv", "weight_scale") + # Shared-expert fusion: redirect the checkpoint shared expert into + # routed-expert slot ``num_local_experts`` (gate->w1, up->w3, + # down->w2) so it loads via the routed expert loader. Runs before the + # stacked/dense mappings so shared_experts.gate_proj/up_proj are not + # captured by the dense gate_up_proj mapping. + if _fuse_shared and ".shared_experts." in name: + sid = self.config.num_local_experts + name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.") + name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.") + name = name.replace(".shared_experts.down_proj.", f".experts.{sid}.w2.") + for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue From 37ce34922f7f5e58241369511130cd99c1c50bfe Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Fri, 26 Jun 2026 16:21:20 +0200 Subject: [PATCH 037/138] [CI] Fix failing CUDA graph capture in Triton MOE (#46735) Signed-off-by: Felix Marty --- .../layers/fused_moe/experts/nvfp4_emulation_moe.py | 5 +++++ vllm/model_executor/layers/fused_moe/experts/triton_moe.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py index f93c67a97dd..cd862cac595 100644 --- a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py @@ -372,6 +372,11 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts): def expects_unquantized_inputs(self) -> bool: return True + @property + def a1_scale(self) -> torch.Tensor: + # Used in experts/triton_moe.py and passed to moe_kernel_quantize_input. + return self.a1_gscale + @staticmethod def supports_lora() -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 8c756e25702..3196667b3f7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -245,7 +245,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): lora_unquantized_hidden_states = hidden_states hidden_states, a1q_scale = moe_kernel_quantize_input( hidden_states, - self.a1_scale or self.a1_gscale, + self.a1_scale, self.quant_dtype, self.per_act_token_quant, self.block_shape, From e71bc6da85577b2057292e60e959ce44af344897 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 26 Jun 2026 23:24:13 +0800 Subject: [PATCH 038/138] [Rust Frontend] Use `oss-harmony` for Harmony output processing (#46799) --- rust/Cargo.lock | 480 ++++++------------------------------------------ rust/Cargo.toml | 2 +- 2 files changed, 53 insertions(+), 429 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 52635c75e8f..e1c051a6fe1 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -31,24 +31,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "aligned" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685" -dependencies = [ - "as-slice", -] - -[[package]] -name = "aligned-vec" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" -dependencies = [ - "equator", -] - [[package]] name = "android_system_properties" version = "0.1.5" @@ -144,12 +126,6 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" -[[package]] -name = "arbitrary" -version = "1.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" - [[package]] name = "arc-swap" version = "1.9.0" @@ -159,17 +135,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "arg_enum_proc_macro" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "arrayref" version = "0.3.9" @@ -182,15 +147,6 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" -[[package]] -name = "as-slice" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516" -dependencies = [ - "stable_deref_trait", -] - [[package]] name = "async-io" version = "2.6.0" @@ -333,49 +289,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "av-scenechange" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f321d77c20e19b92c39e7471cf986812cbb46659d2af674adc4331ef3f18394" -dependencies = [ - "aligned", - "anyhow", - "arg_enum_proc_macro", - "arrayvec", - "log", - "num-rational", - "num-traits", - "pastey", - "rayon", - "thiserror 2.0.18", - "v_frame", - "y4m", -] - -[[package]] -name = "av1-grain" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cfddb07216410377231960af4fcab838eaa12e013417781b78bd95ee22077f8" -dependencies = [ - "anyhow", - "arrayvec", - "log", - "nom 8.0.0", - "num-rational", - "v_frame", -] - -[[package]] -name = "avif-serialize" -version = "0.8.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "375082f007bd67184fb9c0374614b29f9aaa604ec301635f72338bb65386a53d" -dependencies = [ - "arrayvec", -] - [[package]] name = "axum" version = "0.8.8" @@ -484,27 +397,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" -[[package]] -name = "bit_field" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e4b40c7323adcfc0a41c4b88143ed58346ff65a288fc144329c5c45e05d70c6" - [[package]] name = "bitflags" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" -[[package]] -name = "bitstream-io" -version = "4.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7eff00be299a18769011411c9def0d827e8f2d7bf0c3dbf53633147a8867fd1f" -dependencies = [ - "no_std_io2", -] - [[package]] name = "blake3" version = "1.8.5" @@ -539,12 +437,6 @@ dependencies = [ "serde", ] -[[package]] -name = "built" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4ad8f11f288f48ca24471bbd51ac257aaeaaa07adae295591266b792902ae64" - [[package]] name = "bumpalo" version = "3.20.2" @@ -1210,26 +1102,6 @@ dependencies = [ "log", ] -[[package]] -name = "equator" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" -dependencies = [ - "equator-macro", -] - -[[package]] -name = "equator-macro" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "equivalent" version = "1.0.2" @@ -1262,7 +1134,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab" dependencies = [ "futures-core", - "nom 7.1.3", + "nom", "pin-project-lite", ] @@ -1276,21 +1148,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "exr" -version = "1.74.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4300e043a56aa2cb633c01af81ca8f699a321879a7854d3896a0ba89056363be" -dependencies = [ - "bit_field", - "half", - "lebe", - "miniz_oxide", - "rayon-core", - "smallvec", - "zune-inflate", -] - [[package]] name = "fancy-regex" version = "0.13.0" @@ -2022,16 +1879,11 @@ dependencies = [ "bytemuck", "byteorder-lite", "color_quant", - "exr", "gif", "image-webp", "moxcms", "num-traits", "png", - "qoi", - "ravif", - "rayon", - "rgb", "tiff", "zune-core", "zune-jpeg", @@ -2047,12 +1899,6 @@ dependencies = [ "quick-error", ] -[[package]] -name = "imgref" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7c5cedc30da3a610cac6b4ba17597bdf7152cf974e8aab3afb3d54455e371c8" - [[package]] name = "indexmap" version = "1.9.3" @@ -2098,17 +1944,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "interpolate_name" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "ipnet" version = "2.12.0" @@ -2249,28 +2084,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" -[[package]] -name = "lebe" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8" - [[package]] name = "libc" version = "0.2.183" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" -[[package]] -name = "libfuzzer-sys" -version = "0.4.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f12a681b7dd8ce12bff52488013ba614b869148d54dd79836ab85aafdd53f08d" -dependencies = [ - "arbitrary", - "cc", -] - [[package]] name = "libm" version = "0.2.16" @@ -2348,15 +2167,6 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" -[[package]] -name = "loop9" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fae87c125b03c1d2c0150c90365d7d6bcc53fb73a9acaef207d2d065860f062" -dependencies = [ - "imgref", -] - [[package]] name = "macro_rules_attribute" version = "0.2.2" @@ -2456,16 +2266,6 @@ dependencies = [ "rawpointer", ] -[[package]] -name = "maybe-rayon" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ea1f30cedd69f0a2954655f7188c6a834246d2bcf1e315e2ac40c4b24dc9519" -dependencies = [ - "cfg-if", - "rayon", -] - [[package]] name = "memchr" version = "2.8.0" @@ -2637,21 +2437,6 @@ dependencies = [ "rawpointer", ] -[[package]] -name = "new_debug_unreachable" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" - -[[package]] -name = "no_std_io2" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b51ed7824b6e07d354605f4abb3d9d300350701299da96642ee084f5ce631550" -dependencies = [ - "memchr", -] - [[package]] name = "nom" version = "7.1.3" @@ -2662,21 +2447,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "nom" -version = "8.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" -dependencies = [ - "memchr", -] - -[[package]] -name = "noop_proc_macro" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" - [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2686,16 +2456,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "num-bigint" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" -dependencies = [ - "num-integer", - "num-traits", -] - [[package]] name = "num-complex" version = "0.4.6" @@ -2711,17 +2471,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050" -[[package]] -name = "num-derive" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "num-integer" version = "0.1.46" @@ -2731,17 +2480,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "num-traits" version = "0.2.19" @@ -2810,29 +2548,6 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" -[[package]] -name = "openai-harmony" -version = "0.0.8" -source = "git+https://github.com/Inferact/openai-harmony?rev=cfbadbc66f3158692bfeefa961e363aa7a6b9708#cfbadbc66f3158692bfeefa961e363aa7a6b9708" -dependencies = [ - "anyhow", - "base64 0.22.1", - "bstr", - "clap", - "fancy-regex 0.13.0", - "futures", - "image", - "regex", - "reqwest", - "rustc-hash 1.1.0", - "serde", - "serde_json", - "serde_with", - "sha1", - "sha2", - "thiserror 2.0.18", -] - [[package]] name = "openai-protocol" version = "1.6.0" @@ -2912,6 +2627,24 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "oss-harmony" +version = "0.0.11" +source = "git+https://github.com/oss-harmony/harmony?tag=v0.0.11#76e849426cc092f84509e31a17027755f67d662a" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex 0.13.0", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "serde_with", + "sha2", + "thiserror 2.0.18", + "zstd", +] + [[package]] name = "parking" version = "2.2.1" @@ -2947,12 +2680,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pastey" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" - [[package]] name = "pcre2" version = "0.2.11" @@ -3221,25 +2948,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "profiling" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3eb8486b569e12e2c32ad3e204dbaba5e4b5b216e9367044f25f1dba42341773" -dependencies = [ - "profiling-procmacros", -] - -[[package]] -name = "profiling-procmacros" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52717f9a02b6965224f95ca2a81e2e0c5c43baacd28ca057577988930b6c3d5b" -dependencies = [ - "quote", - "syn 2.0.117", -] - [[package]] name = "prometheus-client" version = "0.24.0" @@ -3280,7 +2988,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.10.5", "log", "multimap", "petgraph", @@ -3301,7 +3009,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -3411,15 +3119,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "qoi" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f6d64c71eb498fe9eae14ce4ec935c555749aef511cca85b5568910d6e48001" -dependencies = [ - "bytemuck", -] - [[package]] name = "quick-error" version = "2.0.1" @@ -3506,56 +3205,6 @@ dependencies = [ "getrandom 0.3.4", ] -[[package]] -name = "rav1e" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43b6dd56e85d9483277cde964fd1bdb0428de4fec5ebba7540995639a21cb32b" -dependencies = [ - "aligned-vec", - "arbitrary", - "arg_enum_proc_macro", - "arrayvec", - "av-scenechange", - "av1-grain", - "bitstream-io", - "built", - "cfg-if", - "interpolate_name", - "itertools 0.14.0", - "libc", - "libfuzzer-sys", - "log", - "maybe-rayon", - "new_debug_unreachable", - "noop_proc_macro", - "num-derive", - "num-traits", - "paste", - "profiling", - "rand 0.9.2", - "rand_chacha 0.9.0", - "simd_helpers", - "thiserror 2.0.18", - "v_frame", - "wasm-bindgen", -] - -[[package]] -name = "ravif" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e52310197d971b0f5be7fe6b57530dcd27beb35c1b013f29d66c1ad73fbbcc45" -dependencies = [ - "avif-serialize", - "imgref", - "loop9", - "quick-error", - "rav1e", - "rayon", - "rgb", -] - [[package]] name = "rawpointer" version = "0.2.1" @@ -3680,7 +3329,6 @@ dependencies = [ "base64 0.22.1", "bytes", "encoding_rs", - "futures-channel", "futures-core", "futures-util", "h2", @@ -3726,18 +3374,12 @@ dependencies = [ "futures-core", "futures-timer", "mime", - "nom 7.1.3", + "nom", "pin-project-lite", "reqwest", "thiserror 1.0.69", ] -[[package]] -name = "rgb" -version = "0.8.53" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" - [[package]] name = "ring" version = "0.17.14" @@ -4297,17 +3939,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "sha1" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" -dependencies = [ - "cfg-if", - "cpufeatures 0.2.17", - "digest", -] - [[package]] name = "sha2" version = "0.10.9" @@ -4350,15 +3981,6 @@ version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" -[[package]] -name = "simd_helpers" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95890f873bec569a0362c235787f3aca6e1e887302ba4840839bcc6459c42da6" -dependencies = [ - "quote", -] - [[package]] name = "siphasher" version = "1.0.2" @@ -4405,7 +4027,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" dependencies = [ "base64 0.13.1", - "nom 7.1.3", + "nom", "serde", "unicode-segmentation", ] @@ -5348,17 +4970,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "v_frame" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "666b7727c8875d6ab5db9533418d7c764233ac9c0cff1d469aec8fa127597be2" -dependencies = [ - "aligned-vec", - "num-traits", - "wasm-bindgen", -] - [[package]] name = "validator" version = "0.20.0" @@ -5424,7 +5035,7 @@ dependencies = [ "llm-multimodal", "minijinja", "minijinja-contrib", - "openai-harmony", + "oss-harmony", "paste", "reqwest", "rmp-serde", @@ -6205,12 +5816,6 @@ dependencies = [ "thiserror 2.0.18", ] -[[package]] -name = "y4m" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" - [[package]] name = "yoke" version = "0.8.1" @@ -6345,21 +5950,40 @@ version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "zune-core" version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb8a0807f7c01457d0379ba880ba6322660448ddebc890ce29bb64da71fb40f9" -[[package]] -name = "zune-inflate" -version = "0.2.54" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73ab332fe2f6680068f3582b16a24f90ad7096d5d39b974d1c0aff0125116f02" -dependencies = [ - "simd-adler32", -] - [[package]] name = "zune-jpeg" version = "0.5.15" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 4f6322e7ada..a1f963b9b0f 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -52,7 +52,7 @@ minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builti minijinja-contrib = { version = "2.0", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } ndarray = { version = "0.16.1", features = ["serde"] } -openai-harmony = { git = "https://github.com/Inferact/openai-harmony", rev = "cfbadbc66f3158692bfeefa961e363aa7a6b9708", default-features = false, features = ["native-tls"] } +openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmony/harmony", tag = "v0.0.11", default-features = false } openai-protocol = "1.6.0" parking_lot = "0.12.5" paste = "1.0.15" From 4e07ca2c9284ad0661a44d33e1e8a1c597c48686 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 26 Jun 2026 08:24:33 -0700 Subject: [PATCH 039/138] [Core] Add `VLLM_GPU_SYNC_CHECK` env var (#44800) --- tests/utils_/test_gpu_sync_debug.py | 61 +++++++++ vllm/compilation/compiler_interface.py | 28 +++++ vllm/envs.py | 8 ++ vllm/utils/gpu_sync_debug.py | 165 +++++++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 17 +++ 5 files changed, 279 insertions(+) create mode 100644 tests/utils_/test_gpu_sync_debug.py create mode 100644 vllm/utils/gpu_sync_debug.py diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py new file mode 100644 index 00000000000..ea9b76f34bc --- /dev/null +++ b/tests/utils_/test_gpu_sync_debug.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +import vllm.utils.gpu_sync_debug as gsd +from vllm.utils.gpu_sync_debug import ( + SYNC_ERROR_MESSAGE, + gpu_sync_allowed, + with_gpu_sync_check, +) + +from ..utils import create_new_process_for_each_test + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _no_sync(): + # Pure on-GPU compute, no implicit CPU sync... + x = torch.ones(4, device="cuda") + 1 + # ...plus a sync that we explicitly allow. + with gpu_sync_allowed(): + return x.cpu() + + +def _causes_sync(): + x = torch.ones(4, device="cuda") + # An allowed sync (suppressed)... + with gpu_sync_allowed(): + x.cpu() + # ...then an un-allowed sync that should trip the check. + return x.cpu() + + +@pytest.mark.parametrize("mode", ["warn", "error"]) +@create_new_process_for_each_test() +def test_with_env_set(monkeypatch, mode): + # Env set + gate flipped on: the unguarded sync is detected. + monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", mode) + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + # Guarded syncs always pass. + with_gpu_sync_check(_no_sync)() + + if mode == "error": + # "error" mode turns the stray sync into a RuntimeError. + with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE): + with_gpu_sync_check(_causes_sync)() + else: + # "warn" mode only warns, so the call still succeeds. + with_gpu_sync_check(_causes_sync)() + + +@create_new_process_for_each_test() +def test_without_env_set(monkeypatch): + # Env unset: the decorator is a pass-through, no sync is detected. + monkeypatch.delenv("VLLM_GPU_SYNC_CHECK", raising=False) + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + with_gpu_sync_check(_no_sync)() + with_gpu_sync_check(_causes_sync)() diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 2348ff3191b..742ec55e6ef 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -765,6 +765,34 @@ def set_functorch_config() -> None: setattr(torch._functorch.config, k, v) +def trigger_inductor_lazy_init(device: torch.device | None = None) -> None: + """Eagerly trigger inductor's once-per-process lazy inits (SFDP pattern + matcher, pad_mm, misc patterns). + + These normally fire on the first torch.compile invocation and include + CUDA syncs. If warmup hits the on-disk compile cache, no compile actually + runs so these never fire during warmup, and they'd blow up on the first + real-request cache miss once the sync-check gate is on. + + Private torch API; best-effort. Newer torch versions take an + `input_device` argument and cache per-device, so pass the current CUDA + device to ensure the cache key matches later compile calls. + """ + try: + import inspect + + from torch._inductor.fx_passes.joint_graph import ( + lazy_init as _inductor_lazy_init, + ) + + if inspect.signature(_inductor_lazy_init).parameters: + _inductor_lazy_init(device) + else: + _inductor_lazy_init() + except Exception as e: # noqa: BLE001 + logger.info("Skipping inductor lazy_init pre-trigger: %s", e) + + class EagerAdaptor(CompilerInterface): name = "eager" diff --git a/vllm/envs.py b/vllm/envs.py index 08314a8c88d..ab6184d22dc 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -87,6 +87,7 @@ if TYPE_CHECKING: VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest" VLLM_BATCH_INVARIANT: bool = False VLLM_TRITON_ATTN_USE_TD: bool | None = None + VLLM_GPU_SYNC_CHECK: Literal["warn", "error"] | None = None MAX_JOBS: str | None = None NVCC_THREADS: str | None = None VLLM_USE_PRECOMPILED: bool = False @@ -584,6 +585,13 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_TRITON_ATTN_USE_TD": lambda: {"1": True, "0": False}.get( os.getenv("VLLM_TRITON_ATTN_USE_TD", "").strip() ), + # If set, enable PyTorch's GPU<->CPU synchronization debug mode around + # the worker's `execute_model` and `sample_tokens` calls. Valid values + # are "warn" (print a warning on each sync) or "error" (raise on sync). + # Unset disables the check. See `torch.cuda.set_sync_debug_mode`. + "VLLM_GPU_SYNC_CHECK": env_with_choices( + "VLLM_GPU_SYNC_CHECK", None, ["warn", "error"] + ), # Maximum number of compilation jobs to run in parallel. # By default this is the number of CPUs "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py new file mode 100644 index 00000000000..1e2114f3b5b --- /dev/null +++ b/vllm/utils/gpu_sync_debug.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools +import sys +from contextlib import contextmanager + +import torch + +import vllm.envs as envs +from vllm.platforms import current_platform + +SYNC_ERROR_MESSAGE = ( + "GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()" +) + +_GPU_SYNC_ALLOWED_FIRST_SEEN: set[tuple[str, int]] = set() + +# Global sync-check gate. Off during engine setup (model load, KV cache +# init, warmup/compile) so first-compile and lazy-init syncs pass through; +# flipped on by `enable_gpu_sync_check()` at the end of +# `GPUWorker.compile_or_warm_up_model`, after which `with_gpu_sync_check`- +# decorated functions activate the configured debug mode. +_sync_check_enabled: bool = False + + +def enable_gpu_sync_check() -> None: + """Flip the sync-check gate on. Call once per worker, after warmup / + first-compile is complete. No-op unless `VLLM_GPU_SYNC_CHECK` is set.""" + if envs.VLLM_GPU_SYNC_CHECK is None: + return + global _sync_check_enabled + _sync_check_enabled = True + _install_compile_time_sync_suppressors() + + +_compile_time_suppressors_installed: bool = False + + +def _install_compile_time_sync_suppressors() -> None: + """Wrap torch inductor/aot_autograd compile entry points so the + synchronizing ops those passes perform don't trip the + sync-check mode we set around `execute_model` / `sample_tokens`. + + Warmup-time compiles already run under the gate (before + `enable_gpu_sync_check`), but post-warmup compiles fire inside + `execute_model` and we want to avoid this tripping the sync check. + """ + global _compile_time_suppressors_installed + if _compile_time_suppressors_installed: + return + _compile_time_suppressors_installed = True + + try: # noqa: BLE001 + from torch._inductor.fx_passes import joint_graph as _jg + + _orig_joint = _jg.joint_graph_passes + + @functools.wraps(_orig_joint) + def _wrapped_joint(*args, **kwargs): + prev_mode = torch.cuda.get_sync_debug_mode() + if not prev_mode: + return _orig_joint(*args, **kwargs) + torch.cuda.set_sync_debug_mode(0) + try: + return _orig_joint(*args, **kwargs) + finally: + torch.cuda.set_sync_debug_mode(prev_mode) + + # `compile_fx` does `from .fx_passes.joint_graph import + # joint_graph_passes`, which binds the *function object* at import + # time. Patching just the module attribute won't update that rebind, + # so patch every already-imported reference we can find. Restrict + # the scan to torch's compile-time modules. + import sys as _sys + + setattr(_jg, "joint_graph_passes", _wrapped_joint) # noqa: B010 + for _name, _mod in list(_sys.modules.items()): + if _mod is None: + continue + if not ( + _name.startswith("torch._inductor") + or _name.startswith("torch._functorch") + or _name.startswith("torch._dynamo") + ): + continue + if getattr(_mod, "joint_graph_passes", None) is _orig_joint: + setattr(_mod, "joint_graph_passes", _wrapped_joint) # noqa: B010 + except Exception: # pragma: no cover + pass + + +@contextmanager +def _suppress_gpu_sync_check(prev_mode: int): + torch.cuda.set_sync_debug_mode(0) + try: + yield + finally: + torch.cuda.set_sync_debug_mode(prev_mode) + + +@contextmanager +def _noop_cm(): + yield + + +if current_platform.is_cuda_alike(): + + def gpu_sync_allowed(first_only: bool = False): + """Context manager that suppresses `torch.cuda.set_sync_debug_mode` for the + duration of the `with` block. + + If `first_only` is True, only the first entry from this call site + suppresses the sync check; subsequent entries from the same site are + no-ops so any further GPU syncs will be reported. The "site" is the + caller's (filename, lineno), so different + `with gpu_sync_allowed(first_only=True):` lines track independently. + """ + if envs.VLLM_GPU_SYNC_CHECK is None or torch.compiler.is_compiling(): + return _noop_cm() + prev_mode = torch.cuda.get_sync_debug_mode() + if not prev_mode: + return _noop_cm() + if first_only: + frame = sys._getframe(1) + key = (frame.f_code.co_filename, frame.f_lineno) + if key in _GPU_SYNC_ALLOWED_FIRST_SEEN: + return _noop_cm() + _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key) + return _suppress_gpu_sync_check(prev_mode) + + def with_gpu_sync_check(fn): + """Decorator that enables `torch.cuda.set_sync_debug_mode` around `fn` + when `VLLM_GPU_SYNC_CHECK` is set *and* the gate has been flipped by + `enable_gpu_sync_check()`. Before the gate flips (i.e. during + engine setup / warmup) the decorated function runs as-is. + """ + mode = envs.VLLM_GPU_SYNC_CHECK + if mode is None: + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if not _sync_check_enabled: + return fn(*args, **kwargs) + prev_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode(mode) + try: + return fn(*args, **kwargs) + except RuntimeError as re: + if str(re) == "called a synchronizing CUDA operation": + raise RuntimeError(SYNC_ERROR_MESSAGE) from re + raise re + finally: + torch.cuda.set_sync_debug_mode(prev_mode) + + return wrapper + +else: + # No-op the methods in non-CUDA cases. + + def gpu_sync_allowed(first_only: bool = False): + return _noop_cm() + + def with_gpu_sync_check(fn): + return fn diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 4aa8ca5ca3d..589a16576eb 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -56,6 +56,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.tracing import instrument from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback +from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling from vllm.utils.torch_utils import set_random_seed @@ -759,6 +760,16 @@ class Worker(WorkerBase): # the model initialization and profiling. set_random_seed(self.model_config.seed) + # Eagerly trigger inductor's once-per-process lazy inits during + # warmup (rather than on a later compile cache-miss at runtime). + c_config = self.compilation_config + if c_config.mode != CompilationMode.NONE and c_config.backend == "inductor": + from vllm.compilation.compiler_interface import ( + trigger_inductor_lazy_init, + ) + + trigger_inductor_lazy_init(self.device) + # All warmup is done โ€” start monitoring for unexpected JIT # compilations that would cause latency spikes during inference. from vllm.utils.jit_monitor import activate as activate_jit_monitor @@ -773,6 +784,10 @@ class Worker(WorkerBase): freeze_gc_heap() maybe_attach_gc_debug_callback() + # Warmup / first-compile is done โ€” activate the `VLLM_GPU_SYNC_CHECK` + # gate so subsequent `execute_model` / `sample_tokens` calls enforce it. + enable_gpu_sync_check() + return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, @@ -826,12 +841,14 @@ class Worker(WorkerBase): return self.profiler.annotate_context_manager(annotation) @torch.inference_mode() + @with_gpu_sync_check def sample_tokens( self, grammar_output: "GrammarOutput | None" ) -> ModelRunnerOutput | AsyncModelRunnerOutput: return self.model_runner.sample_tokens(grammar_output) @torch.inference_mode() + @with_gpu_sync_check def execute_model( self, scheduler_output: "SchedulerOutput" ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None: From abc71548ef029132c3316b902207f254a246d593 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Fri, 26 Jun 2026 23:28:49 +0800 Subject: [PATCH 040/138] [CI/Build][CPU] Add test image cache clean-up (#46831) Signed-off-by: jiang1.li --- .../scripts/hardware_ci/run-cpu-test.sh | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 0f0c18b55af..032d8e78333 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -12,6 +12,44 @@ IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}" TIMEOUT_VAL=$1 TEST_COMMAND=$2 +# Disk hygiene knobs. Reclaim space only once the Docker root filesystem crosses +# DISK_USAGE_THRESHOLD percent, and cap the shared BuildKit cache at +# BUILDKIT_CACHE_MAX so subsequent builds keep reusing the hottest layers. +DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-70} +BUILDKIT_CACHE_MAX=${BUILDKIT_CACHE_MAX:-80GB} + +# Reclaim disk only when the host is under pressure. We trim (not purge) the +# shared BuildKit cache so cross-job/cross-agent reuse stays intact, and only +# touch dangling images; other agents' uniquely tagged images are left alone. +prune_if_disk_pressure() { + local docker_root disk_usage + docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true) + if [ -z "$docker_root" ]; then + return 0 + fi + disk_usage=$(df "$docker_root" 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%') + if [ "${disk_usage:-0}" -gt "$DISK_USAGE_THRESHOLD" ]; then + echo "--- :broom: Disk usage ${disk_usage}% exceeds ${DISK_USAGE_THRESHOLD}%, reclaiming space" + docker image prune -f || true + docker builder prune -f --keep-storage="$BUILDKIT_CACHE_MAX" || true + else + echo "Disk usage ${disk_usage:-unknown}% within ${DISK_USAGE_THRESHOLD}% threshold; skipping prune" + fi +} + +# Always drop this agent's image once the job ends (the default builder never +# uses it as a cache source, so removing it costs no rebuild speed), then +# reclaim space if needed. Guard every docker call with `|| true` so the trap +# never overrides the test's exit code. +cleanup() { + docker image rm -f "$IMAGE_NAME" || true + prune_if_disk_pressure +} +trap cleanup EXIT + +# Free space up front so a nearly-full host doesn't fail the build. +prune_if_disk_pressure + # building the docker image echo "--- :docker: Building Docker image" docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu . From 658b54efe419d0e53ec33a8bb7095d8c8b52c741 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 26 Jun 2026 09:36:31 -0700 Subject: [PATCH 041/138] [ModelRunner V2] Update scheduler tests to cover MRV2 paths (#46771) Signed-off-by: Nick Hill --- tests/v1/core/test_scheduler.py | 63 ++++++++++++++++++++++++--------- tests/v1/core/utils.py | 8 ++++- 2 files changed, 53 insertions(+), 18 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index d0168c9a935..dad345c643a 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -6,6 +6,7 @@ from unittest.mock import Mock import pytest import torch +import vllm.envs as envs from vllm.config import ( CacheConfig, ECTransferConfig, @@ -150,7 +151,7 @@ def test_cached_request_data_resumed_all_token_ids_mrv1_only(): """ from vllm.v1.core.kv_cache_manager import KVCacheBlocks - scheduler = create_scheduler() + scheduler = create_scheduler(use_v2_model_runner=False) (req,) = create_requests(num_requests=1, num_tokens=8) req.append_output_token_ids([101, 102, 103]) @@ -1899,11 +1900,14 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role): assert len(scheduler.waiting) == 0 +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.parametrize( "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")] ) -def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role): +def test_kv_connector_handles_preemption( + is_async, use_ec_connector, ec_role, use_v2_model_runner +): """ Test whether scheduler with KVConnector is able to handle unable to allocate (run out of blocks in allocate_slots(). @@ -1924,6 +1928,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role): # encoder connector should not affect test results use_ec_connector=use_ec_connector, ec_role=ec_role, + use_v2_model_runner=use_v2_model_runner, ) # Create two requests. @@ -2034,8 +2039,14 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role): ) assert len(scheduler.running) == 1 assert len(scheduler.waiting) == 0 - assert output.scheduled_cached_reqs.num_reqs == 1 - assert output.scheduled_new_reqs == [] + if use_v2_model_runner: + # V2 emits a resumed (previously preempted) request as a + # NewRequestData rather than a cached request. + assert output.scheduled_cached_reqs.num_reqs == 0 + assert len(output.scheduled_new_reqs) == 1 + else: + assert output.scheduled_cached_reqs.num_reqs == 1 + assert output.scheduled_new_reqs == [] _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT) assert len(scheduler.running) == 1 assert len(scheduler.waiting) == 0 @@ -2155,6 +2166,7 @@ def create_scheduler_with_priority( num_speculative_tokens: int | None = None, use_ec_connector: bool = False, ec_role: str | None = None, + use_v2_model_runner: bool | None = None, ) -> Scheduler: """Create scheduler with priority policy enabled. @@ -2246,7 +2258,7 @@ def create_scheduler_with_priority( ], ) cache_config.num_gpu_blocks = num_blocks - return Scheduler( + scheduler = Scheduler( vllm_config=vllm_config, kv_cache_config=kv_cache_config, log_stats=True, @@ -2254,6 +2266,10 @@ def create_scheduler_with_priority( block_size=block_size, hash_block_size=block_size, ) + if use_v2_model_runner is None: + use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER) + scheduler.use_v2_model_runner = use_v2_model_runner + return scheduler _none_hash_initialized = False @@ -2955,11 +2971,12 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): assert engine_core_output.finish_reason == FinishReason.ERROR +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) @pytest.mark.parametrize( "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")] ) def test_priority_scheduling_preemption_and_resumption_when_out_of_kv( - use_ec_connector, ec_role + use_ec_connector, ec_role, use_v2_model_runner ): """Test that priority scheduling preempts lower priority requests when out of KV cache space.""" @@ -2973,6 +2990,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv( # encoder connector should not affect test results use_ec_connector=use_ec_connector, ec_role=ec_role, + use_v2_model_runner=use_v2_model_runner, ) # Create a request and schedule it @@ -3065,20 +3083,31 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv( output = scheduler.schedule() scheduled_cached_reqs = output.scheduled_cached_reqs - assert len(output.scheduled_new_reqs) == 0 - assert scheduled_cached_reqs.num_reqs == 1 assert len(scheduler.waiting) == 0 assert len(scheduler.running) == 1 - # Preempted request resumed in scheduled_cached_reqs - assert len(scheduled_cached_reqs.resumed_req_ids) == 1 - assert len(scheduled_cached_reqs.all_token_ids) == 1 - assert scheduled_cached_reqs.req_ids[0] == request_low.request_id - assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids - assert request_low.request_id in scheduled_cached_reqs.all_token_ids - # Resumed tokens include 30 prompt tokens and 2 decoded tokens - assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32 - assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100 + if use_v2_model_runner: + # V2 emits the resumed request as a NewRequestData, carrying its full + # token ids in prefill_token_ids (instead of cached all_token_ids). + assert scheduled_cached_reqs.num_reqs == 0 + assert len(output.scheduled_new_reqs) == 1 + new_req = output.scheduled_new_reqs[0] + assert new_req.req_id == request_low.request_id + # Resumed tokens include 30 prompt tokens and 2 decoded tokens. + assert len(new_req.prefill_token_ids) == 32 + assert new_req.prefill_token_ids[31] == 100 + else: + assert len(output.scheduled_new_reqs) == 0 + assert scheduled_cached_reqs.num_reqs == 1 + # Preempted request resumed in scheduled_cached_reqs + assert len(scheduled_cached_reqs.resumed_req_ids) == 1 + assert len(scheduled_cached_reqs.all_token_ids) == 1 + assert scheduled_cached_reqs.req_ids[0] == request_low.request_id + assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids + assert request_low.request_id in scheduled_cached_reqs.all_token_ids + # Resumed tokens include 30 prompt tokens and 2 decoded tokens + assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32 + assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100 @pytest.mark.parametrize( diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 7f34250cb21..2450b23669a 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -3,6 +3,7 @@ import torch +import vllm.envs as envs from tests.v1.kv_connector.unit.utils import MockKVConfig from vllm.config import ( CacheConfig, @@ -58,6 +59,7 @@ def create_scheduler( pipeline_parallel_size: int = 1, use_ec_connector: bool = False, ec_role: str | None = None, + use_v2_model_runner: bool | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -165,13 +167,17 @@ def create_scheduler( cache_config.num_gpu_blocks = num_blocks register_all_kvcache_specs(vllm_config) scheduler_cls = AsyncScheduler if async_scheduling else Scheduler - return scheduler_cls( + scheduler = scheduler_cls( vllm_config=vllm_config, kv_cache_config=kv_cache_config, block_size=block_size, log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), ) + if use_v2_model_runner is None: + use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER) + scheduler.use_v2_model_runner = use_v2_model_runner + return scheduler _none_hash_initialized = False From 3d3b96488f5f7d94b1ab63919e0f1e7922a9ded6 Mon Sep 17 00:00:00 2001 From: Julien Denize <40604584+juliendenize@users.noreply.github.com> Date: Fri, 26 Jun 2026 20:06:31 +0200 Subject: [PATCH 042/138] Migrate Voxtral to mistral-common 1.11.5 audio API (#46705) Signed-off-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> --- examples/generate/multimodal/audio_language_offline.py | 7 ++----- requirements/common.txt | 2 +- requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/nightly-torch.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.txt | 2 +- tests/models/multimodal/generation/test_voxtral.py | 8 +++----- .../multimodal/generation/test_voxtral_realtime.py | 7 +++---- vllm/model_executor/models/voxtral.py | 9 +++++---- vllm/model_executor/models/voxtral_realtime.py | 6 ++---- 12 files changed, 22 insertions(+), 29 deletions(-) diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py index c480f1b4145..12a38cf41cc 100644 --- a/examples/generate/multimodal/audio_language_offline.py +++ b/examples/generate/multimodal/audio_language_offline.py @@ -463,16 +463,15 @@ def run_ultravox(question: str, audio_count: int) -> ModelRequestData: # Voxtral # Make sure to install mistral-common[audio]. def run_voxtral(question: str, audio_count: int) -> ModelRequestData: - from mistral_common.audio import Audio from mistral_common.protocol.instruct.chunk import ( AudioChunk, - RawAudio, TextChunk, ) from mistral_common.protocol.instruct.messages import ( UserMessage, ) from mistral_common.protocol.instruct.request import ChatCompletionRequest + from mistral_common.tokens.tokenizers.audio import Audio from mistral_common.tokens.tokenizers.mistral import MistralTokenizer model_name = "mistralai/Voxtral-Mini-3B-2507" @@ -495,9 +494,7 @@ def run_voxtral(question: str, audio_count: int) -> ModelRequestData: Audio.from_file(str(audio_assets[i].get_local_path()), strict=False) for i in range(audio_count) ] - audio_chunks = [ - AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios - ] + audio_chunks = [AudioChunk.from_audio(audio) for audio in audios] messages = [UserMessage(content=[*audio_chunks, text_chunk])] diff --git a/requirements/common.txt b/requirements/common.txt index a5d74e14e64..1652480c22f 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -33,7 +33,7 @@ partial-json-parser # used for parsing partial JSON outputs jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation pyzmq >= 25.0.0 msgspec -mistral_common[image] >= 1.11.3 +mistral_common[image] >= 1.11.5 opencv-python-headless >= 4.13.0 # required for video 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 diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index a7fc65def8e..03218c75e1f 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -31,7 +31,7 @@ torchaudio==2.11.0 torchvision==0.26.0 transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.11.3 # required for voxtral test +mistral_common[image,audio] >= 1.11.5 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py opencv-python-headless >= 4.13.0 # required for video test diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 76c343b91b1..1a9fe6f16a0 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -398,7 +398,7 @@ mbstrdecoder==1.1.3 # typepy mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.3 +mistral-common==1.11.5 # via # -c requirements/common.txt # -r requirements/test/cuda.in diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index a58e0fa248f..08f721771c8 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -23,7 +23,7 @@ jiwer # required for audio tests timm # required for internvl test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.11.3 # required for voxtral test +mistral_common[image,audio] >= 1.11.5 # required for voxtral test num2words # required for smolvlm test opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 046ca09ff7f..6a38f384f11 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -30,7 +30,7 @@ tblib # for pickling test exceptions timm>=1.0.17 # required for internvl and gemma3n-mm test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio]>=1.11.3 # required for voxtral test +mistral_common[image,audio]>=1.11.5 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py opencv-python-headless>=4.13.0 # required for video test diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 842d2ff3188..726aad9a672 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -499,7 +499,7 @@ mcp==1.27.0 # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.3 +mistral-common==1.11.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 40f23b95d10..2b938e3b583 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -264,7 +264,7 @@ mbstrdecoder==1.1.4 # typepy mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.3 +mistral-common==1.11.5 # via # -c requirements/common.txt # -r requirements/test/xpu.in diff --git a/tests/models/multimodal/generation/test_voxtral.py b/tests/models/multimodal/generation/test_voxtral.py index 82db1dc6812..a6e8f3ff18a 100644 --- a/tests/models/multimodal/generation/test_voxtral.py +++ b/tests/models/multimodal/generation/test_voxtral.py @@ -4,9 +4,9 @@ import json import pytest -from mistral_common.audio import Audio -from mistral_common.protocol.instruct.chunk import AudioChunk, RawAudio, TextChunk +from mistral_common.protocol.instruct.chunk import AudioChunk, TextChunk from mistral_common.protocol.instruct.messages import UserMessage +from mistral_common.tokens.tokenizers.audio import Audio from transformers import VoxtralForConditionalGeneration from vllm.tokenizers.mistral import MistralTokenizer @@ -36,9 +36,7 @@ def _get_prompt(audio_assets: AudioTestAssets, question: str) -> list[int]: Audio.from_file(str(asset.get_local_path()), strict=False) for asset in audio_assets ] - audio_chunks = [ - AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios - ] + audio_chunks = [AudioChunk.from_audio(audio) for audio in audios] messages = [ UserMessage(content=[*audio_chunks, TextChunk(text=question)]).to_openai() diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index ca43e7b51f7..be677ccb570 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -4,12 +4,11 @@ import contextlib import pytest import pytest_asyncio -from mistral_common.audio import Audio -from mistral_common.protocol.instruct.chunk import RawAudio from mistral_common.protocol.transcription.request import ( StreamingMode, TranscriptionRequest, ) +from mistral_common.tokens.tokenizers.audio import Audio from mistral_common.tokens.tokenizers.mistral import MistralTokenizer from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy @@ -101,7 +100,7 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, engine): def from_file(file_path: str): audio = Audio.from_file(file_path, strict=False) req = TranscriptionRequest( - audio=RawAudio.from_audio(audio), + audio=audio.to_base64(audio.format), streaming=StreamingMode.OFFLINE, language=None, ) @@ -156,7 +155,7 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine) req = TranscriptionRequest( streaming=StreamingMode.OFFLINE, - audio=RawAudio.from_audio(audio), + audio=audio.to_base64(audio.format), language=None, ) audio_enc = tokenizer.encode_transcription(req) diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py index baf6d5c7394..f15ec491af1 100644 --- a/vllm/model_executor/models/voxtral.py +++ b/vllm/model_executor/models/voxtral.py @@ -10,11 +10,12 @@ import numpy as np import regex as re import torch import torch.nn as nn -from mistral_common.audio import Audio, mel_filter_bank -from mistral_common.protocol.instruct.chunk import AudioChunk, RawAudio, TextChunk +from mistral_common.audio import mel_filter_bank +from mistral_common.protocol.instruct.chunk import AudioChunk, TextChunk from mistral_common.protocol.instruct.messages import UserMessage from mistral_common.protocol.instruct.request import ChatCompletionRequest from mistral_common.protocol.transcription.request import TranscriptionRequest +from mistral_common.tokens.tokenizers.audio import Audio from transformers import BatchFeature, WhisperConfig from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig @@ -182,7 +183,7 @@ class VoxtralDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralProcessingInfo]): sampling_rate=feature_extractor.sampling_rate, format=format, ) - chunk = AudioChunk(input_audio=RawAudio.from_audio(audio_item)) + chunk = AudioChunk.from_audio(audio_item) audio_chunks.append(chunk) request = ChatCompletionRequest( @@ -462,7 +463,7 @@ class VoxtralForConditionalGeneration( audio = Audio(audio, int(stt_config.sample_rate), format="wav") # lossless req = TranscriptionRequest( model=model_config.model, - audio=RawAudio.from_audio(audio), + audio=audio.to_base64(audio.format), language=language, ) diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 2628e1443e2..8c59532e3b6 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -7,13 +7,11 @@ from collections.abc import AsyncGenerator, Iterable, Iterator, Mapping import numpy as np import torch -from mistral_common.audio import Audio -from mistral_common.protocol.instruct.chunk import RawAudio from mistral_common.protocol.transcription.request import ( StreamingMode, TranscriptionRequest, ) -from mistral_common.tokens.tokenizers.audio import AudioConfig +from mistral_common.tokens.tokenizers.audio import Audio, AudioConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig @@ -477,7 +475,7 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim req = TranscriptionRequest( model=model_config.model, - audio=RawAudio.from_audio(audio), + audio=audio.to_base64(audio.format), language=language, streaming=StreamingMode.OFFLINE, ) From c6554f321ce4c7563290d02eec323f262fc43fef Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 26 Jun 2026 12:32:21 -0600 Subject: [PATCH 043/138] [CPU] Fix macOS/Apple Silicon hang by enabling OpenMP in the build (#46769) Signed-off-by: mgoin Co-authored-by: Claude Opus 4.8 --- .github/actionlint.yaml | 2 ++ .github/workflows/macos-smoke-test.yml | 25 +++++++++++++------ cmake/cpu_extension.cmake | 3 +++ csrc/cpu/cpu_attn_impl.hpp | 6 ++--- csrc/cpu/cpu_fused_moe.cpp | 2 +- csrc/cpu/cpu_types.hpp | 16 ++++++++++++ csrc/cpu/cpu_wna16.cpp | 2 +- csrc/cpu/dnnl_kernels.cpp | 2 +- csrc/cpu/mla_decode.cpp | 2 +- .../installation/cpu.apple.inc.md | 4 +++ 10 files changed, 49 insertions(+), 15 deletions(-) diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml index 940c2885809..082e8a9eb90 100644 --- a/.github/actionlint.yaml +++ b/.github/actionlint.yaml @@ -3,3 +3,5 @@ self-hosted-runner: labels: - vllm-runners + # Not yet in actionlint's known-label set. + - macos-26 diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index 9068ec281b2..eb502578ea4 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -11,7 +11,19 @@ permissions: jobs: macos-m1-smoke-test: - runs-on: macos-latest + # macos-26 (the supported target) is still a preview runner, so gate on GA + # macos-15 and keep macos-26 non-blocking. + strategy: + fail-fast: false + matrix: + include: + - os: macos-15 + required: true + - os: macos-26 + required: false + name: macos-m1-smoke-test (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + continue-on-error: ${{ !matrix.required }} timeout-minutes: 30 steps: @@ -72,14 +84,11 @@ jobs: # Test health endpoint curl -f http://localhost:8000/health - # Test completion - curl -f http://localhost:8000/v1/completions \ + # Long prompt: hits the split-KV path that short prompts skip (#46769). + PAYLOAD=$(python -c "import json; print(json.dumps({'model': 'Qwen/Qwen3-0.6B', 'prompt': 'The quick brown fox jumps over the lazy dog. ' * 24, 'max_tokens': 16}))") + curl -f --max-time 120 http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ - -d '{ - "model": "Qwen/Qwen3-0.6B", - "prompt": "Hello", - "max_tokens": 5 - }' + -d "$PAYLOAD" # Cleanup kill "$SERVER_PID" diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 5c19446601e..9d8796c0d7a 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -24,7 +24,10 @@ set (ENABLE_NUMA TRUE) # Check the compile flags # if(MACOSX_FOUND) + # Apple clang needs -Xpreprocessor to enable OpenMP. No runtime link is + # needed: _C is a dynamic_lookup bundle and resolves libomp from torch. list(APPEND CXX_COMPILE_FLAGS + "-Xpreprocessor" "-fopenmp" "-DVLLM_CPU_EXTENSION") else() list(APPEND CXX_COMPILE_FLAGS diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 7b3757b313d..260ed7cd417 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -124,7 +124,7 @@ struct AttentionMetadata { workitem_group_num(workitem_group_num), reduction_item_num(reduction_item_num), reduction_split_num(reduction_split_num), - thread_num(omp_get_max_threads()), + thread_num(cpu_utils::get_max_threads()), effective_thread_num(thread_num), split_kv_q_token_num_threshold(split_kv_q_token_num_threshold), attention_scratchpad_size_per_thread(0), @@ -405,7 +405,7 @@ class AttentionScheduler { torch::Tensor schedule(const ScheduleInput& input) const { const bool causal = input.causal; const bool is_dynamic_causal = input.dynamic_causal != nullptr; - const int32_t thread_num = omp_get_max_threads(); + const int32_t thread_num = cpu_utils::get_max_threads(); const int64_t cache_size = cpu_utils::get_available_l2_size(); const int32_t max_num_q_per_iter = input.max_num_q_per_iter; const int32_t kv_len_alignment = input.kv_block_alignment; @@ -1423,7 +1423,7 @@ class AttentionMainLoop { public: void operator()(const AttentionInput* input) { - const int thread_num = omp_get_max_threads(); + const int thread_num = cpu_utils::get_max_threads(); TORCH_CHECK_EQ(input->metadata->thread_num, thread_num); std::atomic guard_counter(0); std::atomic* guard_counter_ptr = &guard_counter; diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 35c23df97be..07b0aaf8688 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -267,7 +267,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0); TORCH_CHECK_EQ(output_size_13 / 2, input_size_2); - const int32_t thread_num = omp_get_max_threads(); + const int32_t thread_num = cpu_utils::get_max_threads(); const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( gemm_m_tile_size * input_size_13 * sizeof(scalar_t)); diff --git a/csrc/cpu/cpu_types.hpp b/csrc/cpu/cpu_types.hpp index 744c80c8f53..7b2c3d3b74c 100644 --- a/csrc/cpu/cpu_types.hpp +++ b/csrc/cpu/cpu_types.hpp @@ -25,4 +25,20 @@ #include #endif +#include + +namespace cpu_utils { +// Without OpenMP the omp pragmas compile to serial loops, so report 1: kernels +// that barrier on the thread count would otherwise deadlock. +inline int get_max_threads() { +#ifdef _OPENMP + return omp_get_max_threads(); +#else + TORCH_WARN_ONCE( + "vLLM CPU was built without OpenMP; running single-threaded."); + return 1; +#endif +} +} // namespace cpu_utils + #endif \ No newline at end of file diff --git a/csrc/cpu/cpu_wna16.cpp b/csrc/cpu/cpu_wna16.cpp index 5c6d1ce48a7..ae7aef74c44 100644 --- a/csrc/cpu/cpu_wna16.cpp +++ b/csrc/cpu/cpu_wna16.cpp @@ -155,7 +155,7 @@ void cpu_gemm_wna16_impl( constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; constexpr int32_t n_block_size = 16; static_assert(gemm_n_tile_size % n_block_size == 0); - const int32_t thread_num = omp_get_max_threads(); + const int32_t thread_num = cpu_utils::get_max_threads(); // a simple schedule policy, just to hold more B tiles in L2 and make sure // each thread has tasks diff --git a/csrc/cpu/dnnl_kernels.cpp b/csrc/cpu/dnnl_kernels.cpp index 058fe25b0e2..6dda0929616 100644 --- a/csrc/cpu/dnnl_kernels.cpp +++ b/csrc/cpu/dnnl_kernels.cpp @@ -202,7 +202,7 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output, using cvt_vec_t = typename KernelVecType::cvt_vec_type; constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM; - const int64_t thread_num = omp_get_max_threads(); + const int64_t thread_num = cpu_utils::get_max_threads(); if (num_tokens > thread_num) { #pragma omp parallel for for (int64_t i = 0; i < num_tokens; ++i) { diff --git a/csrc/cpu/mla_decode.cpp b/csrc/cpu/mla_decode.cpp index 3bd0d2e688f..702912a5bcc 100644 --- a/csrc/cpu/mla_decode.cpp +++ b/csrc/cpu/mla_decode.cpp @@ -251,7 +251,7 @@ void mla_decode_kvcache_cpu_impl( constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM; // shared across threads - const int max_threads = omp_get_max_threads(); + const int max_threads = cpu_utils::get_max_threads(); const int acc_out_nbytes = max_threads * num_heads * V_HEAD_DIM * sizeof(float); float* acc_out = static_cast(std::aligned_alloc(64, acc_out_nbytes)); diff --git a/docs/getting_started/installation/cpu.apple.inc.md b/docs/getting_started/installation/cpu.apple.inc.md index e54afc49384..e312964ec8a 100644 --- a/docs/getting_started/installation/cpu.apple.inc.md +++ b/docs/getting_started/installation/cpu.apple.inc.md @@ -15,6 +15,10 @@ Currently the CPU implementation for macOS supports FP32 and FP16 datatypes. - SDK: `XCode 15.4` or later with Command Line Tools - Compiler: `Apple Clang >= 15.0.0` +!!! note + The macOS CPU build is smoke-tested in CI on the latest GA Apple Silicon + runner; other macOS or Apple Clang versions are best-effort. + --8<-- [end:requirements] --8<-- [start:set-up-using-python] From dccb412e2c72a0c147166c14b25c01f045a74163 Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Fri, 26 Jun 2026 15:29:52 -0400 Subject: [PATCH 044/138] [Bugfix][Parser] Pass token IDs to parser.parse() in Responses API and batch serving (#46843) Signed-off-by: Ben Browning --- vllm/entrypoints/openai/chat_completion/batch_serving.py | 1 + vllm/entrypoints/openai/responses/context.py | 1 + vllm/entrypoints/openai/responses/serving.py | 1 + 3 files changed, 3 insertions(+) diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index a0fc8670506..6acc568f492 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -267,6 +267,7 @@ class OpenAIServingChatBatch(OpenAIServingChat): reasoning, content, _ = parser.parse( output.text, request=request, # type: ignore[arg-type] + model_output_token_ids=output.token_ids, ) if not request.include_reasoning: reasoning = None diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 3c9a31a141e..a7cb96f9496 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -346,6 +346,7 @@ class ParsableContext(ConversationContext): completion.text, self.request, enable_auto_tools=self.enable_auto_tools, + model_output_token_ids=completion.token_ids, ) self.response_messages.extend( build_response_output_items( diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index f2e1f8e5d80..6746434a046 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -1063,6 +1063,7 @@ class OpenAIServingResponses(OpenAIServing): final_output.text, request, enable_auto_tools=self.enable_auto_tools, + model_output_token_ids=final_output.token_ids, ) return build_response_output_items( reasoning=reasoning, From 701a23d99f405668158d1395e11c30107dd65b75 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Martinez <45523697+calvarado2004@users.noreply.github.com> Date: Fri, 26 Jun 2026 16:05:04 -0400 Subject: [PATCH 045/138] [Bugfix][Model] Support tensor parallelism for DiffusionGemma (#45719) (#46177) Signed-off-by: Carlos Alvarado Co-authored-by: Claude Co-authored-by: Lucas Wilkinson --- .buildkite/test_areas/lm_eval.yaml | 12 +++++ ...DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml | 10 ++++ tests/evals/gsm8k/configs/models-small-tp.txt | 1 + vllm/model_executor/models/diffusion_gemma.py | 49 +++++++++++++++++-- 4 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml create mode 100644 tests/evals/gsm8k/configs/models-small-tp.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 8063d5e72fd..793b9d8913c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -65,6 +65,18 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt +- label: LM Eval Small Models (2xL4) + key: lm-eval-small-models-tp + timeout_in_minutes: 10 + num_devices: 2 + optional: true + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + autorun_on_main: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt + - label: LM Eval Large Models EP (2xB200) key: lm-eval-large-models-ep-2xb200 timeout_in_minutes: 120 diff --git a/tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml b/tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml new file mode 100644 index 00000000000..3304cbff65c --- /dev/null +++ b/tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml @@ -0,0 +1,10 @@ +model_name: "RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic" +accuracy_threshold: 0.84 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1200 +server_args: >- + --enforce-eager + --max-model-len 4096 + --tensor-parallel-size 2 + --attention-backend TRITON_ATTN diff --git a/tests/evals/gsm8k/configs/models-small-tp.txt b/tests/evals/gsm8k/configs/models-small-tp.txt new file mode 100644 index 00000000000..63bba5bcd1d --- /dev/null +++ b/tests/evals/gsm8k/configs/models-small-tp.txt @@ -0,0 +1 @@ +DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 85e0ef04678..5c4dd8eb554 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -28,6 +28,7 @@ from transformers import AutoModel from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -499,6 +500,13 @@ def _compiled_sample_step( ST: int, # Sampler config entropy_bound: float, + # Tensor-parallel vocab sharding for the self-conditioning matmul. + # ``embed_weight`` is vocab-sharded ([vocab/tp, hidden]) while ``probs`` + # spans the full vocab; [sc_vocab_start, sc_vocab_end) is this rank's slice. + sc_vocab_start: int, + sc_vocab_end: int, + tp_size: int, + tp_group_name: str, ) -> torch.Tensor: """Compiled decode step: temperature โ†’ Gumbel sample โ†’ probs/confidence โ†’ accept/renoise โ†’ convergence, all as vectorized PyTorch ops. @@ -629,7 +637,17 @@ def _compiled_sample_step( # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full # [.., vocab] probs avoids a giant persistent buffer. sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None] - soft_embeds = torch.matmul(probs.to(embed_weight.dtype), embed_weight) * normalizer + # Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor + # parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while + # probs spans the full vocab, so each rank multiplies its local vocab slice + # [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks. + local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype) + soft_embeds = torch.matmul( + local_probs, embed_weight[: sc_vocab_end - sc_vocab_start] + ) + if tp_size > 1: + soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name) + soft_embeds = soft_embeds * normalizer sc_embeds[decode_slots] = soft_embeds * sc_keep # Overwrite canvas with argmax for newly converged denoise requests @@ -843,6 +861,12 @@ class DiffusionGemmaModelState(ModelState): raise ValueError( f"entropy_bound must be a positive float (got {entropy_bound})" ) + # The self-conditioning matmul (probs @ embed_tokens.weight) runs over a + # vocab-parallel embedding shard. Hand the sampler this rank's vocab + # slice and TP group so it can all-reduce the partial products. + embed_tokens = self.model.model.embed_tokens + shard = embed_tokens.shard_indices + tp_group = get_tp_group() return DiffusionSampler( sampler=sampler, diffusion_config=diffusion_config, @@ -852,8 +876,12 @@ class DiffusionGemmaModelState(ModelState): t_max=gen["t_max"], entropy_bound=entropy_bound, confidence_threshold=gen["confidence_threshold"], - embed_weight=self.model.model.embed_tokens.weight, + embed_weight=embed_tokens.weight, normalizer=self.model.model.normalizer, + sc_vocab_start=shard.org_vocab_start_index, + sc_vocab_end=shard.org_vocab_end_index, + tp_size=tp_group.world_size, + tp_group_name=tp_group.unique_name, ), None def apply_staged_writes(self) -> None: @@ -1054,13 +1082,24 @@ class DiffusionSampler: entropy_bound: float, embed_weight: torch.Tensor, normalizer: torch.Tensor, + sc_vocab_start: int = 0, + sc_vocab_end: int | None = None, + tp_size: int = 1, + tp_group_name: str = "", ): self.sampling_states = sampler.sampling_states self.req_states = sampler.req_states # Self-conditioning soft embed = probs @ embed_weight * normalizer, - # computed in the sampler (see _compiled_sample_step). + # computed in the sampler (see _compiled_sample_step). ``embed_weight`` + # is the vocab-parallel shard; [sc_vocab_start, sc_vocab_end) is this + # rank's slice of the full vocab and tp_* drive the cross-rank + # all-reduce. self.embed_weight = embed_weight self.normalizer = normalizer + self.sc_vocab_start = sc_vocab_start + self.sc_vocab_end = sc_vocab_end if sc_vocab_end is not None else vocab_size + self.tp_size = tp_size + self.tp_group_name = tp_group_name self.canvas_length = ( diffusion_config.canvas_length if diffusion_config is not None else 32 ) @@ -1299,6 +1338,10 @@ class DiffusionSampler: CL=self.canvas_length, ST=states.stability_threshold, entropy_bound=self.entropy_bound, + sc_vocab_start=self.sc_vocab_start, + sc_vocab_end=self.sc_vocab_end, + tp_size=self.tp_size, + tp_group_name=self.tp_group_name, ) # --- Logprobs: stash on convergence, return on commit --- From 95e6442a6b6973f827783d162709627304cd13f2 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Fri, 26 Jun 2026 15:19:16 -0500 Subject: [PATCH 046/138] [Hardware][AMD][CI] Fix Kernels Quantization test timeout (#46859) Signed-off-by: Matthew Wong --- .../quantization/test_nvfp4_emulation.py | 57 ++++++++++--------- 1 file changed, 29 insertions(+), 28 deletions(-) diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index f5652af6e92..d2056fe01eb 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -43,6 +43,26 @@ else: return False +MOE_MODEL_CONFIGS = { + "nvidia/Qwen3-30B-A3B-NVFP4": { + "shards": ["model-00001-of-00004.safetensors"], + "expert_prefix": "model.layers.9.mlp.experts.", + # Position of the expert index in the dot-split key. + "expert_idx_pos": 5, + } +} + + +@pytest.fixture(scope="module") +def loaded_model_files(): + return { + model_id: huggingface_hub.snapshot_download( + repo_id=model_id, allow_patterns=config["shards"] + ) + for model_id, config in MOE_MODEL_CONFIGS.items() + } + + class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts): """ Extension of TritonExperts to support emulated NVFP4 MoE experts. @@ -193,17 +213,15 @@ def test_nvfp4_emulation_support_check_rejects_bias_and_lora( not current_platform.is_cuda_alike(), reason="Triton NVFP4 kernel requires CUDA.", ) -def test_triton_dequantize_nvfp4(monkeypatch) -> None: +def test_triton_dequantize_nvfp4(monkeypatch, loaded_model_files) -> None: """Test the Triton dequantization kernel against the CPU reference using real NVFP4 weights from a checkpoint. Tests both 2D (attention projection) and 3D (stacked MoE experts). """ - checkpoint_path = huggingface_hub.snapshot_download( - "nvidia/Qwen3-30B-A3B-NVFP4", - allow_patterns=["model-00001-of-00004.safetensors"], - ) - shard_path = f"{checkpoint_path}/model-00001-of-00004.safetensors" + checkpoint_path = loaded_model_files["nvidia/Qwen3-30B-A3B-NVFP4"] + shards = cast(list[str], MOE_MODEL_CONFIGS["nvidia/Qwen3-30B-A3B-NVFP4"]["shards"]) + shard_path = f"{checkpoint_path}/{shards[0]}" block_size = 16 with safe_open(shard_path, framework="pt", device="cpu") as f: @@ -481,25 +499,8 @@ def test_triton_nvfp4_quant_dequant( print(f" speedup: {speedup:.2f}x") -MOE_MODEL_CONFIGS = { - "nvidia/Qwen3-30B-A3B-NVFP4": { - "shards": ["model-00001-of-00004.safetensors"], - "expert_prefix": "model.layers.9.mlp.experts.", - # Position of the expert index in the dot-split key. - "expert_idx_pos": 5, - }, - "nvidia/Kimi-K2.6-NVFP4": { - "shards": [ - "model-00001-of-00060.safetensors", - "model-00002-of-00060.safetensors", - ], - "expert_prefix": "language_model.model.layers.1.mlp.experts.", - "expert_idx_pos": 6, - }, -} - - def _load_nvfp4_moe_weights( + model_files: dict[str, str], model_id: str, tensor_parallel_size: int, max_experts: int | None = None, @@ -518,10 +519,8 @@ def _load_nvfp4_moe_weights( """ cfg = MOE_MODEL_CONFIGS[model_id] shards = cast(list[str], cfg["shards"]) - checkpoint_path = huggingface_hub.snapshot_download( - model_id, - allow_patterns=shards, - ) + checkpoint_path = model_files[model_id] + expert_prefix = cfg["expert_prefix"] idx_pos = cast(int, cfg["expert_idx_pos"]) @@ -636,6 +635,7 @@ def _load_nvfp4_moe_weights( [pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]], ) def test_nvfp4_moe_correctness( + loaded_model_files, num_tokens: int, top_k: int, model_id: str, @@ -660,6 +660,7 @@ def test_nvfp4_moe_correctness( hidden_dim, intermediate_size, ) = _load_nvfp4_moe_weights( + loaded_model_files, model_id, tensor_parallel_size, max_experts=num_test_experts, From 274325dd43681e1131f22df6d5aad86ac50d9617 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 26 Jun 2026 15:38:38 -0500 Subject: [PATCH 047/138] [ROCm][CI] Remove V1 Sample + Logits from mi250 Queue (#46867) Signed-off-by: Micah Williamson --- .buildkite/test-amd.yaml | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 598940e3e3e..083aa024b2a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -505,26 +505,6 @@ steps: commands: - pytest -v -s v1/attention -- 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 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] From 6e2fb02fe5bcd3990c6ecd2663a5468c27d130f9 Mon Sep 17 00:00:00 2001 From: Charlie Fu Date: Fri, 26 Jun 2026 15:41:49 -0500 Subject: [PATCH 048/138] [ROCm][CI] Fix rlhf_nccl.py on ROCm (#46851) Signed-off-by: charlifu --- examples/rl/rlhf_nccl.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/examples/rl/rlhf_nccl.py b/examples/rl/rlhf_nccl.py index b94d5e4db82..a9e39aaa720 100644 --- a/examples/rl/rlhf_nccl.py +++ b/examples/rl/rlhf_nccl.py @@ -29,6 +29,7 @@ causes unexpected behavior. import os import ray +import torch from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from transformers import AutoModelForCausalLM @@ -39,12 +40,24 @@ from vllm.distributed.weight_transfer.nccl_engine import ( NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine, ) +from vllm.platforms import current_platform from vllm.utils.network_utils import get_ip, get_open_port MODEL_NAME = "facebook/opt-125m" # MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128" +def get_assigned_gpu(): + """This is a temporary workaround for a runtime bug in RCCL on ROCm.""" + if not current_platform.is_rocm(): + return 0 + assigned_gpu = int(ray.get_gpu_ids()[0]) + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + os.environ.pop("HIP_VISIBLE_DEVICES", None) + torch.accelerator.set_device_idx(assigned_gpu) + return assigned_gpu + + class MyLLM(LLM): """Configure the vLLM worker for Ray placement group execution.""" @@ -58,9 +71,11 @@ class TrainModel: """Ray actor that wraps the training model on a dedicated GPU.""" def __init__(self, model_name: str): + assigned_gpu = get_assigned_gpu() + self.model = AutoModelForCausalLM.from_pretrained( model_name, - ).to("cuda:0") + ).to(f"cuda:{assigned_gpu}") self.port = get_open_port() self.master_address = get_ip() From 65e655d2959111d508ad97515c85be0627a7b916 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 26 Jun 2026 14:09:05 -0700 Subject: [PATCH 049/138] [GLM-5] Add DSV3.2/GLM5 to `vllm/models/` (#46808) Signed-off-by: Woosuk Kwon --- vllm/models/deepseek_v32/__init__.py | 22 + vllm/models/deepseek_v32/nvidia/__init__.py | 2 + vllm/models/deepseek_v32/nvidia/attention.py | 423 +++++++++++++++++++ vllm/models/deepseek_v32/nvidia/model.py | 333 +++++++++++++++ vllm/models/deepseek_v32/nvidia/mtp.py | 390 +++++++++++++++++ 5 files changed, 1170 insertions(+) create mode 100644 vllm/models/deepseek_v32/__init__.py create mode 100644 vllm/models/deepseek_v32/nvidia/__init__.py create mode 100644 vllm/models/deepseek_v32/nvidia/attention.py create mode 100644 vllm/models/deepseek_v32/nvidia/model.py create mode 100644 vllm/models/deepseek_v32/nvidia/mtp.py diff --git a/vllm/models/deepseek_v32/__init__.py b/vllm/models/deepseek_v32/__init__.py new file mode 100644 index 00000000000..1b0aa64262f --- /dev/null +++ b/vllm/models/deepseek_v32/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V3.2 (``deepseek_v32``) model โ€” hardware-isolated entry point. + +DeepSeek V3.2 introduced the DeepSeek Sparse Attention (DSA) architecture: +MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA +attend. The same model code serves any DSA checkpoint, including GLM-5.2 +(``glm_moe_dsa``), which reuses this architecture. +""" + +from vllm.platforms import current_platform + +if current_platform.is_rocm() or current_platform.is_xpu(): + raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.") + +from .nvidia.model import DeepseekV32ForCausalLM +from .nvidia.mtp import DeepseekV32MTP + +__all__ = [ + "DeepseekV32ForCausalLM", + "DeepseekV32MTP", +] diff --git a/vllm/models/deepseek_v32/nvidia/__init__.py b/vllm/models/deepseek_v32/nvidia/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py new file mode 100644 index 00000000000..21b0c2c441d --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from transformers import DeepseekV2Config, DeepseekV3Config + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +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.rotary_embedding import get_rope +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.models.deepseek_v2 import ( + DeepSeekV2FusedQkvAProjLinear, + DeepseekV32IndexerCache, + yarn_get_mscale, +) +from vllm.model_executor.models.utils import extract_layer_index +from vllm.utils.torch_utils import is_quantized_kv_cache + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonMetadata, + ) + + +class DeepseekV32Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + prefix: str = "", + ): + super().__init__() + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads + self.head_dim = config.index_head_dim + self.rope_dim = config.qk_rope_head_dim + self.q_lora_rank = q_lora_rank + + # No tensor parallel, just replicated. + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # FP8 wk weights are upcasted to BF16 during loading to keep this fused. + 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 + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 + self.topk_indices_buffer = topk_indices_buffer + + # fp8 naive cache: value in fp8 + fp32 scale per quant_block_size element. + assert cache_config is not None, "DeepSeek V3.2 indexer requires cache_config" + self.k_cache = DeepseekV32IndexerCache( + head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.prefix = prefix + + from vllm.v1.attention.backends.mla.indexer import ( + get_max_prefill_buffer_size, + ) + + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + ) + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> torch.Tensor: + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + + q_pe, q_nope = torch.split( + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + # 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( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) + # RoPE (NeoX) can introduce extra leading dims; reshape back to flat. + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) + k_pe = k_pe.reshape(-1, 1, self.rope_dim) + + q = torch.cat([q_pe, q_nope], dim=-1) + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + + # Only quant q here; k quant is fused with cache insertion. + q = q.view(-1, self.head_dim) + q_fp8, q_scale = per_token_group_quant_fp8( + q, + self.quant_block_size, + column_major_scales=False, + use_ue8m0=self.scale_fmt is not None, + ) + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) + q_scale = q_scale.view(-1, self.n_head, 1) + + weights = ( + weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 + ) + weights = weights.squeeze(-1) + + return self.indexer_op(hidden_states, q_fp8, k, weights) + + +class DeepseekV32Attention(MLAAttention): + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + ) -> None: + quant_config = vllm_config.quant_config + cache_config = vllm_config.cache_config + + hidden_size = config.hidden_size + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + q_lora_rank = config.q_lora_rank + kv_lora_rank = config.kv_lora_rank + num_heads = config.num_attention_heads + + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + num_local_heads = num_heads // tp_size + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + scaling = qk_head_dim**-0.5 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + + # DSA checkpoints may use plain ("default") or yarn-scaled RoPE. + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + if config.rope_parameters["rope_type"] == "deepseek_yarn": + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) + scaling_factor = config.rope_parameters["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + scaling = scaling * mscale * mscale + + # DSA "shared indexer" pattern: only some layers carry an indexer; the + # rest reuse the top-k written by the previous indexer layer into the + # shared topk_indices_buffer. DeepSeek-V3.2 builds it on every layer + # (index_topk_freq defaults to 1); GLM-5.2 uses index_topk_freq=4 so + # only layers [0,1,2,6,10,...] (+ MTP) carry one. + layer_id = extract_layer_index(prefix) + index_topk_freq = getattr(config, "index_topk_freq", 1) + index_topk_pattern = getattr(config, "index_topk_pattern", None) + index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2) + if index_topk_pattern is None: + skip_topk = ( + max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0 + ) + elif 0 <= layer_id < len(index_topk_pattern): + skip_topk = index_topk_pattern[layer_id] == "S" + else: + skip_topk = False + # MTP/nextn layers always build a full indexer (they toggle at runtime). + num_hidden_layers = getattr(config, "num_hidden_layers", None) + is_mtp_layer = num_hidden_layers is not None and layer_id >= num_hidden_layers + + # Build kv_b_proj + indexer first; they are passed to MLAAttention.__init__ + # (which runs nn.Module.__init__ and registers them). + kv_b_proj = ColumnParallelLinear( + kv_lora_rank, + num_heads * (qk_nope_head_dim + v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + indexer = None + if not skip_topk or is_mtp_layer: + indexer = DeepseekV32Indexer( + vllm_config, + config, + hidden_size, + q_lora_rank, + quant_config, + cache_config, + topk_indices_buffer, + prefix=f"{prefix}.indexer", + ) + + # Set up the MLA engine (impl, KV cache, scales, backend, registration, + # and process_weights_after_loading) via the MLAAttention base. + super().__init__( + num_heads=num_local_heads, + scale=scaling, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + kv_b_proj=kv_b_proj, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + use_sparse=True, + indexer=indexer, + topk_indices_buffer=topk_indices_buffer, + ) + + self.num_local_heads = num_local_heads + self.qk_head_dim = qk_head_dim + self.indexer = indexer + # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0 + # computes the top-k, steps 1+ set this True to reuse it. + self.skip_topk = False + # Whether the paged KV cache must be viewed as fp8 before the attention + # (per-tensor fp8; the fp8_ds_mla layout is read as uint8). + self._fp8_kv_needs_view = ( + is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype != "fp8_ds_mla" + ) + # Whether the backend takes an fp8-quantized query (FlashInfer sparse) + # vs the (ql_nope, q_pe) tuple (FlashMLA sparse). + self._use_concat_quant = ( + is_quantized_kv_cache(self.kv_cache_dtype) + and self.impl.supports_quant_query_input + ) + + # Remaining MLA projections (registered on this module). + self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( + hidden_size, + [q_lora_rank, kv_lora_rank + qk_rope_head_dim], + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + ) + self.q_a_layernorm = RMSNorm(q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + q_lora_rank, + num_heads * qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps) + self.o_proj = RowParallelLinear( + num_heads * v_head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + # Lightning indexer uses its own RoPE; interleave maps to non-NeoX. + self.indexer_rope_emb = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=not getattr(config, "indexer_rope_interleave", False), + ) + + def forward( # type: ignore[override] + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] + q_c, kv_lora = qkv_lora.split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + ) + q_c = self.q_a_layernorm(q_c) + q = self.q_b_proj(q_c)[0] + + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv_c_normed = self.kv_a_layernorm(kv_c) + + q = q.view(-1, self.num_local_heads, self.qk_head_dim) + k_pe = k_pe.unsqueeze(1) + q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( + positions, q[..., self.qk_nope_head_dim :], k_pe + ) + + num_tokens = hidden_states.shape[0] + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_nope = q_nope.transpose(0, 1) # (N, B, P) + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) # (B, N, L) + + # Lightning indexer writes the top-k indices into the shared buffer. + # "Shared" layers (indexer is None) reuse the top-k from the previous + # indexer layer already sitting in the buffer. + if self.indexer is not None and not self.skip_topk: + self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) # type: ignore[operator] + + attn_latent = torch.empty( + (num_tokens, self.num_local_heads, self.kv_lora_rank), + dtype=q.dtype, + device=q.device, + ) + self._sparse_attention(kv_c_normed, k_pe, ql_nope, q_pe, attn_latent) + + # V up-projection + output projection are metadata-independent GEMMs and + # stay captured. + output = torch.empty( + (num_tokens, self.num_local_heads * self.v_head_dim), + dtype=q.dtype, + device=q.device, + ) + self._v_up_proj(attn_latent, out=output) + return self.o_proj(output)[0] + + @eager_break_during_capture + def _sparse_attention( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + attn_latent: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + attn_metadata: MLACommonMetadata | None + if isinstance(attn_metadata_raw, dict): + attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment] + elif isinstance(attn_metadata_raw, list): + # Speculative decoding: [0] is the base-model metadata dict. + attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment] + else: + attn_metadata = attn_metadata_raw + + slot_mapping = forward_context.slot_mapping + assert isinstance(slot_mapping, dict) + self.impl.do_kv_cache_update( # type: ignore[attr-defined] + kv_c_normed, + k_pe, + self.kv_cache, + slot_mapping.get(self.layer_name), + self.kv_cache_dtype, + self._k_scale, + ) + + if attn_metadata is None: + # Profile / warmup: zero-fill for DP+EP determinism. + attn_latent.zero_() + return + + num_actual = attn_metadata.num_actual_tokens + kv_cache = self.kv_cache + if self._fp8_kv_needs_view: + kv_cache = kv_cache.view(torch.float8_e4m3fn) + + ql_nope = ql_nope[:num_actual] + q_pe = q_pe[:num_actual] + # FlashInfer sparse takes a single fp8-quantized query; FlashMLA sparse + # takes the (ql_nope, q_pe) tuple and concatenates internally. + mqa_q: torch.Tensor | tuple[torch.Tensor, torch.Tensor] + if self._use_concat_quant: + mqa_q = self._decode_concat_quant_fp8_op(ql_nope, q_pe, self._q_scale) + else: + mqa_q = (ql_nope, q_pe) + + attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] + attn_latent[:num_actual] = attn_out.view( + num_actual, self.num_local_heads, self.kv_lora_rank + ) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py new file mode 100644 index 00000000000..dd9e1d65ead --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import torch + +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV2ForCausalLM, + DeepseekV2MLP, + DeepseekV2MoE, + _try_load_fp8_indexer_wk, + get_spec_layer_idx_from_weight_name, +) +from vllm.model_executor.models.utils import ( + PPMissingLayer, + get_pp_missing_layer_names, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, +) +from vllm.sequence import IntermediateTensors + +from .attention import DeepseekV32Attention + + +class DeepseekV32DecoderLayer(torch.nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config=None, + topk_indices_buffer: torch.Tensor | None = None, + ) -> None: + super().__init__() + + if config is None: + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + + self.hidden_size = config.hidden_size + moe_layer_freq = getattr(config, "moe_layer_freq", 1) + layer_idx = int(prefix.split(sep=".")[-1]) + self.layer_idx = layer_idx + self.use_mha = False + + self.self_attn = DeepseekV32Attention( + vllm_config=vllm_config, + config=config, + prefix=f"{prefix}.self_attn", + topk_indices_buffer=topk_indices_buffer, + ) + + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % moe_layer_freq == 0 + ): + self.mlp = DeepseekV2MoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = DeepseekV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class DeepseekV32Model(torch.nn.Module): + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + from vllm.platforms import current_platform + + self.device = current_platform.device_type + + self.vocab_size = config.vocab_size + # DSA is always sparse (has index_topk); allocate the shared top-k + # buffer the indexer writes and the sparse MLA backend reads. + self.is_v32 = True + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV32DecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + topk_indices_buffer=topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + self.aux_hidden_state_layers = tuple[int, ...]() + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + assert input_ids is not None + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + aux_hidden_states = [] + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + if idx in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states + residual) + hidden_states, residual = layer(positions, hidden_states, residual) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states, _ = self.norm(hidden_states, residual) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # DSA-only: MLA (fused_qkv_a_proj) + the fused indexer wk/weights_proj + + # routed experts. No MHA (qkv_proj) or ROCm shared-expert-fusion paths. + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj", "q_a_proj", 0), + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + num_redundant_experts=self.num_redundant_experts, + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # MTP / nextn layers are loaded by the MTP model, not here. + if get_spec_layer_idx_from_weight_name(self.config, name) is not None: + continue + if _try_load_fp8_indexer_wk( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + pp_missing_layer_names, + ): + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Experts are handled below; skip here before the name rewrite. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + if ( + param_name == "fused_qkv_a_proj" + ) and name_mapped not in params_dict: + continue + name = name_mapped + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment] + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + else: + if is_expert_weight: + continue + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment] + if name is None: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): + """DSA causal LM โ€” DeepSeek V2/V3 orchestration with the DSA backbone. + + Serves DeepSeek V3.2 and any architecture reusing DSA (e.g. GLM-5.2). + """ + + model_cls = DeepseekV32Model + + def set_moe_parameters(self): + # Same as the base, but keyed on the MoE block type rather than the + # decoder-layer type (DeepseekV32DecoderLayer is a plain nn.Module). + self.expert_weights = [] + self.num_expert_groups = getattr(self.config, "n_group", 1) + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if isinstance(layer.mlp, DeepseekV2MoE): + example_moe = layer.mlp + self.moe_mlp_layers.append(layer.mlp) + self.moe_layers.append(layer.mlp.experts) + self.extract_moe_parameters(example_moe) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py new file mode 100644 index 00000000000..482ebecc526 --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import typing +from collections.abc import Callable, Iterable + +import torch +import torch.nn as nn + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.deepseek_mtp import SharedHead +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV2MixtureOfExperts, + DeepseekV2MoE, + _try_load_fp8_indexer_wk, + get_spec_layer_idx_from_weight_name, +) +from vllm.model_executor.models.utils import ( + get_pp_missing_layer_names, + maybe_prefix, +) +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + +from .model import DeepseekV32DecoderLayer + + +class DeepseekV32MultiTokenPredictorLayer(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=current_platform.device_type, + ) + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + self.mtp_block = DeepseekV32DecoderLayer( + vllm_config, + prefix, + config=config, + topk_indices_buffer=topk_indices_buffer, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + inputs_embeds = self.enorm(inputs_embeds) + previous_hidden_states = self.hnorm(previous_hidden_states) + hidden_states = self.eh_proj( + torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + ) + hidden_states, residual = self.mtp_block( + positions=positions, hidden_states=hidden_states, residual=None + ) + # Return the pre-final-norm recycle hidden (re-fed as the next spec + # step's previous_hidden_states); shared_head norm is applied in + # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract. + return residual + hidden_states + + +class DeepseekV32MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.layers = torch.nn.ModuleDict( + { + str(idx): DeepseekV32MultiTokenPredictorLayer( + vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def set_skip_topk(self, skip: bool): + # index_share_for_mtp_iteration: step 0 computes top-k, steps 1+ reuse. + for layer in self.layers.values(): + self_attn = getattr(layer.mtp_block, "self_attn", None) + if self_attn is not None and hasattr(self_attn, "skip_topk"): + self_attn.skip_topk = skip + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + return self.logits_processor( + mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) + ) + + +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.model = DeepseekV32MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.set_moe_parameters() + + def set_moe_parameters(self): + self.expert_weights = [] + self.num_moe_layers = self.config.num_nextn_predict_layers + self.num_expert_groups = self.config.n_group + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in self.model.layers.values(): + mlp = layer.mtp_block.mlp + if isinstance(mlp, DeepseekV2MoE): + example_moe = mlp + self.moe_mlp_layers.append(mlp) + self.moe_layers.append(mlp.experts) + self.extract_moe_parameters(example_moe) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + rocm_aiter_moe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj", "q_a_proj", 0), + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + + ( + self.config.n_shared_experts + if rocm_aiter_moe_shared_expert_enabled + else 0 + ), + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + is_fusion_moe_shared_experts_layer = ( + 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, + pp_missing_layer_names, + ): + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if ("mlp.experts." in name) and name not in params_dict: + continue + if is_fusion_moe_shared_experts_layer: + continue + name_mapped = name.replace(weight_name, param_name) + if ( + param_name == "fused_qkv_a_proj" + ) and name_mapped not in params_dict: + continue + else: + name = name_mapped + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + num_chunks = 1 + if is_fusion_moe_shared_experts_layer: + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1 + split_dim = ( + 1 + if ("down_proj.weight" in name and loaded_weight.ndim > 1) + else 0 + ) + total = loaded_weight.shape[split_dim] + assert total % num_chunks == 0 + chunk_size = total // num_chunks + + for j in range(num_chunks): + chunk_name = name + weight_to_load = loaded_weight + if is_fusion_moe_shared_experts_layer: + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size) + if loaded_weight.ndim == 1: + weight_to_load = loaded_weight[chunk_slice] + elif split_dim == 0: + weight_to_load = loaded_weight[chunk_slice, :] + else: + weight_to_load = loaded_weight[:, chunk_slice] + chunk_name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts + j}", + ) + + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment] + if weight_name not in chunk_name: + continue + is_expert_weight = True + name_mapped = chunk_name.replace(weight_name, param_name) + param = params_dict[name_mapped] + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + weight_to_load, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + if not is_fusion_moe_shared_experts_layer: + name = name_mapped + else: + loaded_params.add(name_mapped) + break + else: + if is_expert_weight: + continue + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment] + if name is None: + continue + if ( + spec_layer != self.model.mtp_start_layer_idx + and ".layers" not in name + ): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + if not is_fusion_moe_shared_experts_layer: + loaded_params.add(name) + + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint." + ) + return loaded_params From c40d307731b82a9d472001c5feff18c24797d9df Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Fri, 26 Jun 2026 23:16:39 +0200 Subject: [PATCH 050/138] [Core] Remove FlashAttention block size restriction for hybrid models (#36701) Signed-off-by: Thomas Parnell Co-authored-by: Claude Opus 4.6 --- vllm/v1/attention/backends/flash_attn.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 6aeb7b024b4..75231bafeed 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -43,7 +43,6 @@ if is_flash_attn_varlen_func_available(): import vllm.envs as envs from vllm.config import ( VllmConfig, - get_current_vllm_config, get_current_vllm_config_or_none, get_layers_from_vllm_config, ) @@ -75,22 +74,6 @@ class FlashAttentionBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - vllm_config = get_current_vllm_config() - model_config = vllm_config.model_config - cache_config = vllm_config.cache_config - if ( - model_config - and model_config.is_hybrid - and ( - cache_config.mamba_ssm_cache_dtype == "float32" - or cache_config.mamba_cache_dtype == "float32" - ) - ): - # NOTE(tdoublep): while in principle, FA supports - # MultipleOf(16), these are the block sizes that do not - # suffer from the NaN propagation problem described here: - # https://github.com/Dao-AILab/flash-attention/issues/1974 - return [16, 32, 64] return [MultipleOf(16)] forward_includes_kv_cache_update: bool = False From 77f8796d164ae938072f78561aa72da14990419c Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Fri, 26 Jun 2026 17:18:47 -0400 Subject: [PATCH 051/138] [Frontend][Gpt-oss] Use `process_eos()` to flush Harmony Parser outputs. (#46437) Signed-off-by: Yifan Zong --- tests/parser/test_harmony.py | 31 ++++++++- vllm/parser/harmony.py | 121 ++++++++++++++++++----------------- 2 files changed, 91 insertions(+), 61 deletions(-) diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py index e6646eb763e..f9ca0b7b329 100644 --- a/tests/parser/test_harmony.py +++ b/tests/parser/test_harmony.py @@ -131,6 +131,31 @@ def tool_call_entries(delta_message) -> list[tuple[int, str | None, str | None]] ] +class TestFlush: + def test_flush(self, harmony_parser): + harmony_parser.process_chunk( + encode_output("<|channel|>analysis<|message|>Think") + ) + + flushed = harmony_parser.flush() + + assert flushed is not None + assert flushed.channel == "analysis" + assert flushed.recipient is None + assert flushed.delta == "" + assert flushed.completed_message is not None + assert get_text(flushed.completed_message) == "Think" + assert harmony_parser._parser is None + + def test_flush_resets_after_eos_error(self, harmony_parser): + harmony_parser.process_chunk(encode_output("<|channel|>analysis")) + + flushed = harmony_parser.flush() + + assert flushed is None + assert harmony_parser._parser is None + + class TestParse: # Rendered conversation outputs. @@ -339,6 +364,7 @@ class TestParse: assert reasoning is None assert content == "I'm in the middle of answering" assert tool_calls is None + assert harmony_parser._parser is None def test_interrupted_reasoning_first_message(self, harmony_parser, chat_request): reasoning, content, tool_calls = harmony_parser.parse( @@ -352,6 +378,7 @@ class TestParse: assert reasoning == "I'm in the middle of thinking" assert content is None assert tool_calls is None + assert harmony_parser._parser is None def test_truncated_output(self, harmony_parser, chat_request): reasoning, content, tool_calls = harmony_parser.parse( @@ -367,6 +394,7 @@ class TestParse: assert reasoning == "I'm thinking." assert content == "I'm in the middle of answering" assert tool_calls is None + assert harmony_parser._parser is None @pytest.mark.parametrize( ("harmony_str", "expected_content"), @@ -435,7 +463,7 @@ class TestParseDelta: "<|end|><|start|>assistant<|channel|>final<|message|>Answer" ), request=chat_request, - finished=False, + finished=True, ) assert first_delta is not None @@ -444,6 +472,7 @@ class TestParseDelta: assert second_delta is not None assert second_delta.content == "Answer" assert second_delta.reasoning is None + assert parser._parser is None def test_multi_token(self, gpt_oss_tokenizer, chat_request): parser = HarmonyParser(gpt_oss_tokenizer) diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index ff022a00eb7..4919e3da7eb 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -3,12 +3,15 @@ from __future__ import annotations +import contextlib import json from collections.abc import Sequence from dataclasses import dataclass from enum import Enum, auto from typing import TYPE_CHECKING, NamedTuple +from openai_harmony import HarmonyError + from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.engine.protocol import ( @@ -28,8 +31,7 @@ from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser if TYPE_CHECKING: - from openai_harmony import Message, Role - from openai_harmony import StreamState as HarmonyStreamState + from openai_harmony import Message, StreamableParser class _SegmentType(Enum): @@ -82,33 +84,46 @@ class HarmonyParser(DelegatingParser): f"got {self.tool_parser.__class__.__name__}." ) - self._harmony_parser = get_streamable_parser_for_assistant() + self._parser: StreamableParser | None = None self._next_tool_call_index = 0 self._num_processed_messages = 0 @property - def state(self) -> HarmonyStreamState: - return self._harmony_parser.state + def _harmony_parser(self) -> StreamableParser: + """Lazily initializes the Harmony parser.""" + if self._parser is None: + self._parser = get_streamable_parser_for_assistant() + return self._parser - @property - def current_role(self) -> Role | None: - return self._harmony_parser.current_role + def _poll_completed_message(self) -> Message | None: + messages = self._harmony_parser.messages + if len(messages) <= self._num_processed_messages: + return None + msg = messages[self._num_processed_messages] + self._num_processed_messages += 1 + return msg - @property - def current_channel(self) -> str | None: - return self._harmony_parser.current_channel + def flush(self) -> Segment | None: + msg = None + with contextlib.suppress(HarmonyError): + self._harmony_parser.process_eos() + # TODO: Consider reraising - @property - def current_recipient(self) -> str | None: - return self._harmony_parser.current_recipient + msg = self._poll_completed_message() - @property - def current_content(self) -> str: - return self._harmony_parser.current_content + # Reset to the initial assistant-parser state for the next turn. + self._parser = None + self._num_processed_messages = 0 - @property - def current_content_type(self) -> str | None: - return self._harmony_parser.current_content_type + if msg is None: + return None + + return Segment( + channel=msg.channel, + recipient=msg.recipient, + delta="", + completed_message=msg, + ) def parse( self, @@ -123,24 +138,32 @@ class HarmonyParser(DelegatingParser): Callers must decide whether to surface them. """ result = self.process_chunk(model_output_token_ids) + flushed_segment = self.flush() + if flushed_segment is not None: + result.segments.append(flushed_segment) reasoning_parts: list[str] = [] content_parts: list[str] = [] tool_calls: list[FunctionCall] = [] - def _append_parsed_message( - channel: str | None, - recipient: str | None, - text: str, - content_type: str | None = None, - ) -> None: - segment_type = _SegmentType.from_channel_and_recipient(channel, recipient) + for segment in result.segments: + msg = segment.completed_message + if msg is None: + continue + if msg.author.role != "assistant" or not msg.content: + continue + text = msg.content[0].text + segment_type = _SegmentType.from_channel_and_recipient( + msg.channel, msg.recipient + ) match segment_type: case _SegmentType.REASONING if self.reasoning_parser and text: reasoning_parts.append(text) case _SegmentType.CONTENT if text: content_parts.append(text) case _SegmentType.TOOL if self.tool_parser: + recipient = msg.recipient + content_type = msg.content_type assert recipient is not None if content_type is not None and "json" not in content_type: arguments = text @@ -156,31 +179,6 @@ class HarmonyParser(DelegatingParser): ) ) - for segment in result.segments: - msg = segment.completed_message - if msg is None: - continue - if msg.author.role != "assistant" or not msg.content: - continue - _append_parsed_message( - channel=msg.channel, - recipient=msg.recipient, - text=msg.content[0].text, - content_type=msg.content_type, - ) - - if ( - self.current_channel is not None - or self.current_recipient is not None - or self.current_content - ): - _append_parsed_message( - channel=self.current_channel, - recipient=self.current_recipient, - text=self.current_content, - content_type=self.current_content_type, - ) - reasoning = "\n".join(reasoning_parts) or None content = "\n".join(content_parts) or None return reasoning, content, tool_calls or None @@ -194,8 +192,12 @@ class HarmonyParser(DelegatingParser): *, finished: bool, ) -> DeltaMessage | None: - prev_recipient = self.current_recipient + prev_recipient = self._harmony_parser.current_recipient result = self.process_chunk(delta_token_ids) + if finished: + flushed_segment = self.flush() + if flushed_segment is not None: + result.segments.append(flushed_segment) combined_content = "" combined_reasoning = "" tool_messages: list[DeltaToolCall] = [] @@ -248,6 +250,9 @@ class HarmonyParser(DelegatingParser): ) ) + if finished: + self._next_tool_call_index = 0 + if not combined_content and not combined_reasoning and not tool_messages: return None @@ -268,14 +273,10 @@ class HarmonyParser(DelegatingParser): reasoning_token_count = 0 for token_id in token_ids: self._harmony_parser.process(token_id) - channel = self.current_channel - recipient = self.current_recipient + channel = self._harmony_parser.current_channel + recipient = self._harmony_parser.current_recipient delta = self._harmony_parser.last_content_delta or "" - completed_message = None - _messages = self._harmony_parser.messages - if len(_messages) > self._num_processed_messages: - completed_message = _messages[self._num_processed_messages] - self._num_processed_messages += 1 + completed_message = self._poll_completed_message() if channel == "analysis" or ( channel == "commentary" and recipient is not None From 75fdcc82a5a5ee859e46b489f78630ab61ed40b7 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Fri, 26 Jun 2026 14:48:53 -0700 Subject: [PATCH 052/138] [CI] Add @ivanium to CODEOWNERS for KV-cache/offload areas (#46873) Signed-off-by: Yifan Qiao --- .github/CODEOWNERS | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 15bd35f80e4..8ca6fc22d64 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,7 +3,7 @@ # This lists cover the "core" components of vLLM that require careful review /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng -/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi +/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi @ivanium /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @@ -11,7 +11,7 @@ /vllm/model_executor/layers/mamba @tdoublep @tomeras91 /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn -/vllm/model_executor/layers/batch_invariant.py @yewentao256 +/vllm/model_executor/layers/batch_invariant.py @yewentao256 /vllm/ir @ProExpertProg /vllm/kernels/ @ProExpertProg @tjtanaa /vllm/kernels/helion @ProExpertProg @zou3519 @@ -23,7 +23,7 @@ # Any change to the VllmConfig changes can have a large user-facing impact, # so spam a lot of people /vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg -/vllm/config/cache.py @heheda12345 +/vllm/config/cache.py @heheda12345 @ivanium # Config utils /vllm/config/utils.py @hmellor @@ -67,16 +67,17 @@ /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety @vadiklyutiy /vllm/v1/attention/backends/triton_attn.py @tdoublep /vllm/v1/attention/backends/gdn_attn.py @ZJY0516 @vadiklyutiy -/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery +/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett -/vllm/v1/kv_cache_interface.py @heheda12345 +/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium /vllm/v1/kv_offload @ApostaC @orozery +/vllm/v1/simple_kv_offload @ivanium /vllm/v1/engine @njhill /vllm/v1/executor @njhill /vllm/v1/worker @njhill -/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche +/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche @ivanium # Model runner V2 /vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256 @@ -103,13 +104,14 @@ /tests/test_inputs.py @DarkLight1337 @ywang96 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm /tests/v1/structured_output @mgoin @russellb @aarnphm -/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery +/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/weight_loading @mgoin @youkaichao @yewentao256 /tests/lora @jeejeelee /tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 /tests/v1/kv_connector/nixl_integration @NickLucche -/tests/v1/kv_connector @ApostaC @orozery +/tests/v1/kv_connector @ApostaC @orozery @ivanium /tests/v1/kv_offload @ApostaC @orozery +/tests/v1/simple_kv_offload @ivanium /tests/v1/determinism @yewentao256 /tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning /tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning From 2ff76a5e856e385f72aa49cadc4d0a724d1f7da8 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 26 Jun 2026 16:58:40 -0500 Subject: [PATCH 053/138] [ROCm][Bugfix] Pass num_kv_splits to aiter mla_reduce_v1 (#46760) Signed-off-by: Rohan Potdar Co-authored-by: Claude Opus 4.8 (1M context) --- vllm/v1/attention/backends/mla/rocm_aiter_mla.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index b172370a9f9..41924889d57 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -828,6 +828,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): attn_metadata.fp8_prefill_reduce_final_map, attn_metadata.fp8_prefill_reduce_partial_map, tile_q, + # num_kv_splits added by ROCm/aiter#3391; 0 selects the kernel + # default max(cu_num, 0) == cu_num, matching pre-#3391 behavior. + 0, out_3d, final_lse, ) From d8eb734d94fea27cfcc95a22f3cc2a249e0996c7 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 27 Jun 2026 00:16:05 +0100 Subject: [PATCH 054/138] Fix Transformers backend FP8 MoE and remove some boilerplate (#46820) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/AXK1.py | 2 - vllm/model_executor/models/afmoe.py | 36 ++++---- vllm/model_executor/models/deepseek_mtp.py | 1 - vllm/model_executor/models/deepseek_v2.py | 2 - vllm/model_executor/models/ernie45_moe.py | 2 - vllm/model_executor/models/gemma4.py | 1 - vllm/model_executor/models/gemma4_mm.py | 1 - vllm/model_executor/models/glm4_moe.py | 2 - vllm/model_executor/models/glm4_moe_lite.py | 2 - .../models/glm4_moe_lite_mtp.py | 2 - vllm/model_executor/models/glm4_moe_mtp.py | 2 - vllm/model_executor/models/glm_ocr_mtp.py | 1 - vllm/model_executor/models/hunyuan_v1.py | 1 - vllm/model_executor/models/hy_v3.py | 5 +- vllm/model_executor/models/interfaces.py | 83 +++++++++---------- vllm/model_executor/models/interns1_pro.py | 2 - vllm/model_executor/models/lfm2_moe.py | 1 - vllm/model_executor/models/llama4.py | 2 - vllm/model_executor/models/mellum.py | 2 - vllm/model_executor/models/mixtral.py | 1 - vllm/model_executor/models/nemotron_h.py | 1 - vllm/model_executor/models/openpangu.py | 1 - vllm/model_executor/models/param2moe.py | 3 +- vllm/model_executor/models/qwen3_5.py | 2 - vllm/model_executor/models/qwen3_moe.py | 1 - vllm/model_executor/models/qwen3_next.py | 2 - vllm/model_executor/models/qwen3_vl_moe.py | 2 - vllm/model_executor/models/step3p5.py | 26 +----- .../model_executor/models/transformers/moe.py | 47 +++-------- 29 files changed, 80 insertions(+), 156 deletions(-) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index d526f57d3d9..a465c6b5632 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -1088,8 +1088,6 @@ class AXK1ForCausalLM( self.set_moe_parameters() def set_moe_parameters(self): - self.expert_weights = [] - self.num_expert_groups = getattr(self.config, "n_group", 1) self.moe_layers = [] diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 1564ee733f6..369b7c3b3ad 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -42,6 +42,7 @@ from vllm.model_executor.model_loader.weight_utils import ( ) from vllm.model_executor.models.interfaces import ( EagleModelMixin, + MixtureOfExperts, SupportsEagle3, SupportsLoRA, SupportsPP, @@ -595,7 +596,9 @@ class AfmoeModel(nn.Module, EagleModelMixin): return loaded_params -class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA): +class AfmoeForCausalLM( + nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA, MixtureOfExperts +): packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -635,8 +638,6 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA): self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors ) - self.expert_weights = [] - # Set MoE hyperparameters self.num_moe_layers = config.num_hidden_layers - config.num_dense_layers self.num_expert_groups = config.n_group @@ -663,21 +664,24 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA): self.num_shared_experts = example_moe.n_shared_experts self.num_redundant_experts = example_moe.n_redundant_experts - def set_eplb_state( + def update_physical_experts_metadata( self, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, ) -> None: - for layer_idx, layer in enumerate(self.moe_layers): - # Register the expert weights. - self.expert_weights.append(layer.get_expert_weights()) - layer.set_eplb_state( - moe_layer_idx=layer_idx, - expert_load_view=expert_load_view, - logical_to_physical_map=logical_to_physical_map, - logical_replica_count=logical_replica_count, - ) + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if layer.moe_enabled: + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 88f33ac021b..f73d9f9c3ef 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -217,7 +217,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): self.set_moe_parameters() def set_moe_parameters(self): - self.expert_weights = [] self.num_moe_layers = self.config.num_nextn_predict_layers self.num_expert_groups = self.config.n_group diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 8d20e0b5c68..814118f8a79 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1684,8 +1684,6 @@ class DeepseekV2ForCausalLM( self.set_moe_parameters() def set_moe_parameters(self): - self.expert_weights = [] - self.num_expert_groups = getattr(self.config, "n_group", 1) self.moe_layers = [] diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index e1b9ca9bf57..c2d9f92a666 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -656,8 +656,6 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe self.model.make_empty_intermediate_tensors ) - self.expert_weights = [] - # Set MoE hyperparameters moe_layers_indices = [ i diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 03e67c4ada7..9cf86a5ba83 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -1562,7 +1562,6 @@ class Gemma4ForCausalLM( ) # --- MixtureOfExperts protocol --- - self.expert_weights: list[list[torch.Tensor]] = [] self.moe_layers: list[nn.Module] = [] example_moe: Gemma4MoE | None = None diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index bad7e061cc3..30c379d86c1 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -1114,7 +1114,6 @@ class Gemma4ForConditionalGeneration( ) # --- MixtureOfExperts delegation to language_model --- - self.expert_weights = self.language_model.expert_weights self.moe_layers = self.language_model.moe_layers self.num_moe_layers = self.language_model.num_moe_layers self.num_logical_experts = self.language_model.num_logical_experts diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 98cc9a50adc..8226b65c45c 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -655,8 +655,6 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, Glm4MixtureOfExper self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors ) - self.expert_weights = [] - # Set MoE hyperparameters self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace self.num_expert_groups = config.n_group diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index b4d0fe96680..432fa5e6fa0 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -573,8 +573,6 @@ class Glm4MoeLiteForCausalLM( self.set_moe_parameters() def set_moe_parameters(self): - self.expert_weights = [] - self.num_expert_groups = getattr(self.config, "n_group", 1) self.moe_layers = [] diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index 4813af5f030..222705c14ee 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -209,8 +209,6 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - self.expert_weights = [] - # Set MoE hyperparameters self.num_moe_layers = self.config.num_nextn_predict_layers self.num_expert_groups = self.config.n_group diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index d87ad268285..b255b67d885 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -195,8 +195,6 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - self.expert_weights = [] - # Set MoE hyperparameters self.num_moe_layers = self.config.num_nextn_predict_layers self.num_expert_groups = self.config.n_group diff --git a/vllm/model_executor/models/glm_ocr_mtp.py b/vllm/model_executor/models/glm_ocr_mtp.py index 3d283c101ca..9b2369f93d3 100644 --- a/vllm/model_executor/models/glm_ocr_mtp.py +++ b/vllm/model_executor/models/glm_ocr_mtp.py @@ -134,7 +134,6 @@ class GlmOcrMTP(nn.Module, SupportsPP): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - self.expert_weights = [] self.num_layers = self.config.num_nextn_predict_layers for layer in self.model.layers.values(): assert isinstance(layer, GlmOcrMultiTokenPredictorLayer) diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index ec3cfbd017b..4f70a966289 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -991,7 +991,6 @@ class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts): super().__init__(vllm_config=vllm_config, prefix=prefix) # Set MoE hyperparameters - self.expert_weights = [] self.num_expert_groups = 1 self.moe_layers = [] example_layer = None diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 7653cddd6c7..a4b52e20bda 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -68,7 +68,7 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.hy_v3 import HYV3Config -from .interfaces import SupportsLoRA, SupportsPP +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -392,7 +392,7 @@ class HYV3DecoderLayer(nn.Module): @support_torch_compile -class HYV3Model(nn.Module): +class HYV3Model(nn.Module, MixtureOfExperts): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -429,7 +429,6 @@ class HYV3Model(nn.Module): ) # Set MoE hyperparameters - self.expert_weights = [] self.num_expert_groups = 1 self.moe_layers = [] example_layer = None diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index ad3d01ae9ec..29603318c15 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -29,36 +29,34 @@ from torch import Tensor from transformers.models.whisper.tokenization_whisper import LANGUAGES from typing_extensions import Self, TypeIs -from vllm.config import ModelConfig, SpeechToTextConfig, SpeechToTextParams -from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.tasks import ScoreType from vllm.utils.collection_utils import common_prefix from vllm.utils.func_utils import supports_kw -from .interfaces_base import VllmModel - if TYPE_CHECKING: - from vllm.config import VllmConfig + from vllm.config import ( + ModelConfig, + SpeechToTextConfig, + SpeechToTextParams, + VllmConfig, + ) + from vllm.inputs import PromptType, TokensPrompt from vllm.lora.model_manager import LoRAModelManager + from vllm.model_executor.layers.fused_moe import MoERunner + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc + from vllm.model_executor.models.interfaces_base import VllmModel from vllm.model_executor.models.utils import WeightsMapper from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.multimodal.registry import _ProcessorFactories from vllm.sequence import IntermediateTensors + from vllm.tasks import ScoreType from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, EncoderCudaGraphConfig, EncoderCudaGraphReplayBuffers, EncoderItemSpec, ) -else: - VllmConfig = object - WeightsMapper = object - MultiModalFeatureSpec = object - _ProcessorFactories = object - IntermediateTensors = object logger = init_logger(__name__) @@ -89,7 +87,7 @@ def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor: # Cache results of `SupportsMultiModal.get_language_model` -_language_model_by_module = dict[nn.Module, VllmModel]() +_language_model_by_module = dict[nn.Module, "VllmModel"]() @runtime_checkable @@ -123,7 +121,7 @@ class SupportsMultiModal(Protocol): in their raw form and not input embeddings. """ - _processor_factory: ClassVar[_ProcessorFactories] + _processor_factory: ClassVar["_ProcessorFactories"] """ Set internally by `MultiModalRegistry.register_processor`. """ @@ -175,7 +173,7 @@ class SupportsMultiModal(Protocol): self._has_oov_mm_tokens, ) - def get_language_model(self) -> VllmModel: + def get_language_model(self) -> "VllmModel": """ Returns the underlying language model used for text generation. @@ -216,7 +214,7 @@ class SupportsMultiModal(Protocol): @contextmanager def _mark_language_model( self, - vllm_config: VllmConfig, + vllm_config: "VllmConfig", *, targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, ): @@ -251,7 +249,7 @@ class SupportsMultiModal(Protocol): @contextmanager def _mark_tower_model( self, - vllm_config: VllmConfig, + vllm_config: "VllmConfig", modalities: set[str] | str, *, targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None, @@ -295,7 +293,7 @@ class SupportsMultiModal(Protocol): @contextmanager def _mark_composite_model( self, - vllm_config: VllmConfig, + vllm_config: "VllmConfig", *, language_targets: type[nn.Module] | tuple[type[nn.Module], ...], tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]], @@ -513,7 +511,7 @@ class SupportsScoreTemplate(Protocol): ... @classmethod - def post_process_tokens(cls, prompt: TokensPrompt) -> None: + def post_process_tokens(cls, prompt: "TokensPrompt") -> None: """ Perform architecture-specific manipulations on the input tokens. """ @@ -633,7 +631,7 @@ class SupportsPP(Protocol): batch_size: int, dtype: torch.dtype, device: torch.device, - ) -> IntermediateTensors: + ) -> "IntermediateTensors": """Called when PP rank > 0 for profiling purposes.""" ... @@ -642,8 +640,8 @@ class SupportsPP(Protocol): input_ids: Tensor | None, positions: Tensor, *, - intermediate_tensors: IntermediateTensors | None, - ) -> IntermediateTensors | None: + intermediate_tensors: "IntermediateTensors | None", + ) -> "IntermediateTensors | None": """ Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when PP rank > 0. @@ -665,15 +663,15 @@ class _SupportsPPType(Protocol): batch_size: int, dtype: torch.dtype, device: torch.device, - ) -> IntermediateTensors: ... + ) -> "IntermediateTensors": ... def forward( self, input_ids: Tensor | None, positions: Tensor, *, - intermediate_tensors: IntermediateTensors | None, - ) -> Tensor | IntermediateTensors: ... + intermediate_tensors: "IntermediateTensors | None", + ) -> "Tensor | IntermediateTensors": ... @overload @@ -803,7 +801,7 @@ class IsHybrid(Protocol): @classmethod def get_mamba_state_shape_from_config( cls, - vllm_config: VllmConfig, + vllm_config: "VllmConfig", ) -> tuple[tuple[int, int], tuple[int, int, int]]: """Calculate shapes for Mamba's convolutional and state caches. @@ -818,7 +816,7 @@ class IsHybrid(Protocol): ... @classmethod - def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]: + def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: """Calculate copy-function callables for each Mamba state. Returns: @@ -883,7 +881,7 @@ class MixtureOfExperts(Protocol): num_redundant_experts: int """Number of redundant experts in this model.""" - moe_layers: Iterable[nn.Module] + moe_layers: Iterable["MoERunner"] """List of MoE layers in this model.""" def set_eplb_state( @@ -908,6 +906,7 @@ class MixtureOfExperts(Protocol): logical_to_physical_map: Mapping from logical to physical experts. logical_replica_count: Count of replicas for each logical expert. """ + self.expert_weights = [] for layer_idx, layer in enumerate(self.moe_layers): # Register the expert weights. self.expert_weights.append(layer.get_expert_weights()) @@ -982,7 +981,7 @@ def supports_mamba_prefix_caching( class SupportsCrossEncoding(Protocol): """The interface required for all models that support cross encoding.""" - score_type: ClassVar[ScoreType] = "cross-encoder" + score_type: ClassVar["ScoreType"] = "cross-encoder" @runtime_checkable @@ -994,13 +993,13 @@ class SupportsLateInteraction(Protocol): MaxSim (max over document tokens, sum over query tokens). """ - score_type: ClassVar[ScoreType] = "late-interaction" + score_type: ClassVar["ScoreType"] = "late-interaction" class SupportsQuant: """The interface required for all models that support quantization.""" - hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None + hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None quant_config: QuantizationConfig | None = None @@ -1054,8 +1053,8 @@ class SupportsRealtime(Protocol): cls, audio_stream: AsyncGenerator[np.ndarray, None], input_stream: asyncio.Queue[list[int]], - model_config: ModelConfig, - ) -> AsyncGenerator[PromptType, None]: ... + model_config: "ModelConfig", + ) -> AsyncGenerator["PromptType", None]: ... @overload @@ -1124,8 +1123,8 @@ class SupportsTranscription(Protocol): @classmethod def get_generation_prompt( cls, - stt_params: SpeechToTextParams, - ) -> PromptType: + stt_params: "SpeechToTextParams", + ) -> "PromptType": """Get the prompt for the ASR model. The model has control over the construction, as long as it returns a valid PromptType.""" @@ -1163,8 +1162,8 @@ class SupportsTranscription(Protocol): @classmethod def get_speech_to_text_config( - cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"] - ) -> SpeechToTextConfig: + cls, model_config: "ModelConfig", task_type: Literal["transcribe", "translate"] + ) -> "SpeechToTextConfig": """Get the speech to text config for the ASR model.""" ... @@ -1172,8 +1171,8 @@ class SupportsTranscription(Protocol): def get_num_audio_tokens( cls, audio_duration_s: float, - stt_config: SpeechToTextConfig, - model_config: ModelConfig, + stt_config: "SpeechToTextConfig", + model_config: "ModelConfig", ) -> int | None: """ Map from audio duration to number of audio tokens produced by the ASR @@ -1202,8 +1201,8 @@ class SupportsTranscription(Protocol): def get_language_detection_prompt( cls, audio: np.ndarray, - stt_config: SpeechToTextConfig, - ) -> PromptType: + stt_config: "SpeechToTextConfig", + ) -> "PromptType": """Return a prompt that triggers language detection. Only needs to be implemented when diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 36f669179c5..c04b4729454 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -513,8 +513,6 @@ class InternS1ProMoeMixtureOfExperts(MixtureOfExperts): moe.experts.update_expert_map() def set_moe_parameters(self): - self.expert_weights = [] - self.moe_layers = [] example_moe = None for layer in self.language_model.model.layers: diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 9ca7fb7aaa6..94f7f4e2890 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -692,7 +692,6 @@ class Lfm2MoeForCausalLM( ) # Set MoE hyperparameters - self.expert_weights = [] self.moe_layers = [] example_layer = None diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 9222405ba6d..71df54a4241 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -729,8 +729,6 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts): self.set_moe_parameters() def set_moe_parameters(self): - self.expert_weights = [] - self.moe_layers = [] example_moe = None for layer in self.model.layers: diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py index bdbf0df7fd1..c20fa00e3d6 100644 --- a/vllm/model_executor/models/mellum.py +++ b/vllm/model_executor/models/mellum.py @@ -227,8 +227,6 @@ class MellumForCausalLM(Qwen3MoeForCausalLM): self.model.make_empty_intermediate_tensors ) - self.expert_weights = [] - self.moe_layers = [] example_layer = None for layer in self.model.layers: diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 53c1c87cfce..57eb820ad93 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -512,7 +512,6 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts): self.model.make_empty_intermediate_tensors ) - self.expert_weights = [] self.moe_layers = [] example_moe = None diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 769504c0d0f..bd5cd358d8a 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -886,7 +886,6 @@ class NemotronHForCausalLM( # Set MoE hyperparameters if self.model.has_moe: - self.expert_weights = [] self.num_expert_groups = config.n_group self.moe_layers = [] diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 8432566a150..91120840bdf 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -1289,7 +1289,6 @@ class OpenPanguMoEModel(OpenPanguModelBase, MixtureOfExperts): config = vllm_config.model_config.hf_config # Set MoE hyperparameters - self.expert_weights = [] self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace self.num_expert_groups = 1 diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index e8ea2dbc0e6..3386f9545fa 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -751,7 +751,7 @@ class Param2MoEMixtureOfExperts(MixtureOfExperts): logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, ) -> None: - self.expert_weights.clear() + self.expert_weights = [] for layer_idx, layer in enumerate(self.moe_layers): if hasattr(layer, "get_expert_weights"): self.expert_weights.append(layer.get_expert_weights()) @@ -832,7 +832,6 @@ class Param2MoEForCausalLM( self.model.make_empty_intermediate_tensors ) - self.expert_weights: list[torch.Tensor] = [] self.num_moe_layers: int = 0 self.moe_layers: list = [] self.moe_mlp_layers: list = [] diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index b00b4958681..480ef3678c8 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -765,8 +765,6 @@ class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts): moe.experts.update_expert_map() def set_moe_parameters(self): - self.expert_weights = [] - self.moe_layers = [] example_moe = None for layer in self.language_model.model.layers: diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 6980184cc8a..b7a78acc0ec 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -703,7 +703,6 @@ class Qwen3MoeForCausalLM( ) # Set MoE hyperparameters - self.expert_weights = [] self.moe_layers = [] example_layer = None diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 2c7667a416f..74c2b1e44ad 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -763,8 +763,6 @@ class QwenNextMixtureOfExperts(MixtureOfExperts): moe.experts.update_expert_map() def set_moe_parameters(self): - self.expert_weights = [] - self.moe_layers = [] example_moe = None for layer in self.model.layers: diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 298863209d5..5291874dd5c 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -373,8 +373,6 @@ class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts): moe.experts.update_expert_map() def set_moe_parameters(self): - self.expert_weights = [] - self.moe_layers = [] example_moe = None for layer in self.language_model.model.layers: diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 7a60946ba57..f8bd529e276 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -919,17 +919,17 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts): ) # Set MoE hyperparameters - self.moe_layers: list[FusedMoEBlock] = [] + self.moe_layers: list[MoERunner] = [] + example_layer: FusedMoEBlock | None = None for layer in self.model.layers: if isinstance(layer, PPMissingLayer): continue assert isinstance(layer, Step3p5DecoderLayer) if hasattr(layer, "moe") and isinstance(layer.moe, FusedMoEBlock): - self.moe_layers.append(layer.moe) + example_layer = layer.moe + self.moe_layers.append(layer.moe.experts) - self.expert_weights = [] assert len(self.moe_layers) > 0, "No MoE layers found in the model." - example_layer = self.moe_layers[0] self.num_moe_layers = len(self.moe_layers) self.num_expert_groups = 1 self.num_shared_experts = 0 @@ -959,24 +959,6 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_tokens(input_ids) - def set_eplb_state( - self, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - for layer_idx, layer in enumerate(self.moe_layers): - experts = layer.experts - assert isinstance(experts, MoERunner) - # Register the expert weights. - self.expert_weights.append(experts.get_expert_weights()) - experts.set_eplb_state( - moe_layer_idx=layer_idx, - expert_load_view=expert_load_view, - logical_to_physical_map=logical_to_physical_map, - logical_replica_count=logical_replica_count, - ) - def update_physical_experts_metadata( self, num_physical_experts: int, diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 372c5b1ec12..3e04aaa0748 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -28,14 +28,9 @@ from vllm.config.utils import getattr_iter from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - MoERunner, - fused_moe_make_expert_params_mapping, -) +from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts from vllm.model_executor.models.interfaces import MixtureOfExperts from vllm.model_executor.models.utils import maybe_prefix -from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op from .utils import log_replacement @@ -52,7 +47,7 @@ class TransformersMoEState: # --8<-- [start:transformers_fused_moe] @PluggableLayer.register("transformers_fused_moe") -class TransformersFusedMoE(MoERunner): +class TransformersMoERunner(MoERunner): """Custom FusedMoE for the Transformers modeling backend.""" # --8<-- [end:transformers_fused_moe] @@ -93,7 +88,7 @@ class TransformersFusedMoE(MoERunner): return self.routed_experts.load_weights(weights) -def transformers_moe_forward( +def _transformers_moe_forward( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, @@ -106,7 +101,7 @@ def transformers_moe_forward( return self._forward_super(hidden_states, topk_weights) -def transformers_moe_forward_fake( +def _transformers_moe_forward_fake( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, @@ -117,10 +112,9 @@ def transformers_moe_forward_fake( direct_register_custom_op( op_name="transformers_moe_forward", - op_func=transformers_moe_forward, + op_func=_transformers_moe_forward, mutates_args=["hidden_states"], - fake_impl=transformers_moe_forward_fake, - dispatch_key=current_platform.dispatch_key, + fake_impl=_transformers_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -131,20 +125,6 @@ class MoEMixin(MixtureOfExperts): # Skip MixtureOfExperts.__init__ and call the next class in MRO super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix) - def set_eplb_state( - self, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ): - for moe_layer_idx, mlp_layer in enumerate(self.mlp_moe_layers): - mlp_layer.experts.set_eplb_state( - moe_layer_idx=moe_layer_idx, - expert_load_view=expert_load_view, - logical_to_physical_map=logical_to_physical_map, - logical_replica_count=logical_replica_count, - ) - def update_physical_experts_metadata( self, num_physical_experts: int, @@ -154,7 +134,7 @@ class MoEMixin(MixtureOfExperts): self.num_physical_experts = num_physical_experts self.num_local_physical_experts = num_local_physical_experts self.num_redundant_experts = num_physical_experts - self.num_logical_experts - for mlp in self.mlp_moe_layers: + for mlp in self.mlp_layers: mlp.n_local_physical_experts = num_local_physical_experts mlp.n_physical_experts = num_physical_experts mlp.n_redundant_experts = self.num_redundant_experts @@ -185,7 +165,7 @@ class MoEMixin(MixtureOfExperts): num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts for gate_proj, down_proj, up_proj in ckpt_names: expert_mapping.extend( - fused_moe_make_expert_params_mapping( + RoutedExperts.make_expert_params_mapping( self, ckpt_gate_proj_name=gate_proj, ckpt_down_proj_name=down_proj, @@ -248,10 +228,8 @@ class MoEMixin(MixtureOfExperts): # MixtureOfExperts mixin settings ep_size = get_ep_group().world_size - self.mlp_moe_layers = [] # Used for MixtureOfExperts methods + self.mlp_layers = [] # Used for MixtureOfExperts methods self.moe_layers = [] - self.expert_weights = [] - self.num_moe_layers = 0 self.num_expert_groups = 1 if num_expert_group is None else num_expert_group self.num_logical_experts = num_experts self.num_physical_experts = num_experts + num_redundant_experts @@ -335,19 +313,18 @@ class MoEMixin(MixtureOfExperts): custom_routing_function, moe_state=moe_state, ), - runner_cls=TransformersFusedMoE, + runner_cls=TransformersMoERunner, runner_args={"moe_state": moe_state}, ) mlp.experts = fused_experts log_replacement(qual_name, experts, fused_experts) # Update MixtureOfExperts mixin state - self.mlp_moe_layers.append(mlp) + self.mlp_layers.append(mlp) self.moe_layers.append(fused_experts) - self.expert_weights.append(fused_experts.get_expert_weights()) - self.num_moe_layers += 1 else: _recursive_replace(child_module, prefix=qual_name) _recursive_replace(self.model, prefix="model") + self.num_moe_layers = len(self.moe_layers) # Continue with the replacement of layers in Base super().recursive_replace() From b94f212e37f4ddf4b5e1cc96cd87217f36e3ec0c Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 26 Jun 2026 16:32:45 -0700 Subject: [PATCH 055/138] [ModelRunner V2] Deduplicate ModelState init logic (#46776) Signed-off-by: Nick Hill --- vllm/model_executor/models/diffusion_gemma.py | 28 +------------------ vllm/v1/worker/gpu/model_states/default.py | 26 +---------------- .../gpu/model_states/encoder_decoder.py | 20 +------------ vllm/v1/worker/gpu/model_states/interface.py | 27 ++++++++++++++---- 4 files changed, 25 insertions(+), 76 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 5c4dd8eb554..e28e2720a7f 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -776,33 +776,7 @@ class DiffusionGemmaModelState(ModelState): encoder_cache: Any, device: torch.device, ) -> None: - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - self.scheduler_config = vllm_config.scheduler_config - self.model = model - self.device = device - - self.supports_mm_inputs = encoder_cache is not None - self.max_num_reqs = self.scheduler_config.max_num_seqs - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens - self.max_model_len = self.model_config.max_model_len - self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() - self.dtype = self.model_config.dtype - - if self.supports_mm_inputs: - from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache - from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner - - assert isinstance(encoder_cache, EncoderCache) - self.encoder_cache = encoder_cache - self.encoder_runner = EncoderRunner( - model=self.model, - max_num_tokens=self.max_num_tokens, - hidden_size=self.inputs_embeds_size, - encoder_cache=encoder_cache, - dtype=self.dtype, - device=self.device, - ) + super().__init__(vllm_config, model, encoder_cache, device) # Per-step MM data produced by get_mm_embeddings and consumed by # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 05ff8278864..22e6aa00bc9 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -12,7 +12,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache -from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner from vllm.v1.worker.gpu.mm.rope import get_rope_state from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner @@ -28,30 +27,7 @@ class DefaultModelState(ModelState): encoder_cache: EncoderCache | None, device: torch.device, ): - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - self.scheduler_config = vllm_config.scheduler_config - self.model = model - self.device = device - - self.supports_mm_inputs = encoder_cache is not None - self.max_model_len = self.model_config.max_model_len - self.max_num_reqs = self.scheduler_config.max_num_seqs - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens - self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() - self.dtype = self.model_config.dtype - - if self.supports_mm_inputs: - assert encoder_cache is not None - self.encoder_cache = encoder_cache - self.encoder_runner = EncoderRunner( - model=self.model, - max_num_tokens=self.max_num_tokens, - hidden_size=self.inputs_embeds_size, - encoder_cache=encoder_cache, - dtype=self.dtype, - device=self.device, - ) + super().__init__(vllm_config, model, encoder_cache, device) self.rope_state = get_rope_state( self.model_config, diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index 6ad9a448bee..889e624623d 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -13,7 +13,6 @@ from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache -from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner from vllm.v1.worker.gpu.model_states.interface import ( ModelSpecificAttnMetadata, ModelState, @@ -53,25 +52,8 @@ class EncoderDecoderModelState(ModelState): encoder_cache: EncoderCache | None, device: torch.device, ) -> None: - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - self.scheduler_config = vllm_config.scheduler_config - self.model = model - self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens - self.max_model_len = self.model_config.max_model_len - self.device = device - assert encoder_cache is not None - self.encoder_cache = encoder_cache - self.encoder_runner = EncoderRunner( - model=self.model, - max_num_tokens=self.max_num_tokens, - hidden_size=self.model_config.get_inputs_embeds_size(), - encoder_cache=self.encoder_cache, - dtype=self.model_config.dtype, - device=self.device, - ) + super().__init__(vllm_config, model, encoder_cache, device) self.max_encoder_len = getattr( self.model_config.hf_config, diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 882b38073c4..a4c436a423b 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -37,7 +37,6 @@ class ModelSpecificAttnMetadata: class ModelState(ABC): - @abstractmethod def __init__( self, vllm_config: VllmConfig, @@ -45,11 +44,29 @@ class ModelState(ABC): encoder_cache: EncoderCache | None, device: torch.device, ) -> None: - raise NotImplementedError + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.scheduler_config = vllm_config.scheduler_config + self.model = model + self.device = device - model: nn.Module - # Set by mm-capable states; used by the default gather_mm_embeddings(). - encoder_runner: EncoderRunner + self.max_model_len = self.model_config.max_model_len + self.max_num_reqs = self.scheduler_config.max_num_seqs + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens + self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() + self.dtype = self.model_config.dtype + + self.supports_mm_inputs = encoder_cache is not None + if encoder_cache is not None: + self.encoder_cache = encoder_cache + self.encoder_runner = EncoderRunner( + model=self.model, + max_num_tokens=self.max_num_tokens, + hidden_size=self.inputs_embeds_size, + encoder_cache=encoder_cache, + dtype=self.dtype, + device=self.device, + ) def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]: from vllm.model_executor.models.interfaces import ( From 1d41009e81eb6493f2c19e9d2a0d472564764e62 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 26 Jun 2026 16:34:21 -0700 Subject: [PATCH 056/138] [ModelRunner V2] Fix cross-attention block table sizing (#46753) Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/model_runner.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index de8cd476cd4..cb46ffc3dc0 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -409,10 +409,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): block_table_max_model_len = self.max_model_len if self.is_encoder_decoder: - # Cross-attention block tables need to index encoder tokens - # (e.g., Whisper ~1500), which can exceed decoder max_model_len. + # Cross-attention block tables need to index encoder tokens, which + # can exceed the decoder's max_model_len. block_table_max_model_len = max( block_table_max_model_len, + self.scheduler_config.max_num_encoder_input_tokens, getattr(self.model_config.hf_config, "max_source_positions", 0), ) From 3f674774970225a4aaaa7272a56b3a4c4604eaa7 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 27 Jun 2026 00:56:39 +0100 Subject: [PATCH 057/138] [CI] Don't try and download files that we already know don't exist (#46854) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/transformers_utils/test_repo_utils.py | 29 +++++++++++++++++++++ vllm/transformers_utils/config.py | 6 ++--- vllm/transformers_utils/repo_utils.py | 29 ++++++++++++++++----- 3 files changed, 55 insertions(+), 9 deletions(-) diff --git a/tests/transformers_utils/test_repo_utils.py b/tests/transformers_utils/test_repo_utils.py index 6da4256cba9..36d0acccd6b 100644 --- a/tests/transformers_utils/test_repo_utils.py +++ b/tests/transformers_utils/test_repo_utils.py @@ -7,9 +7,11 @@ from pathlib import Path from unittest.mock import MagicMock, call, patch import pytest +from huggingface_hub import _CACHED_NO_EXIST from vllm.transformers_utils.repo_utils import ( any_pattern_in_repo_files, + get_hf_file_to_dict, is_mistral_model_repo, list_filtered_repo_files, ) @@ -115,6 +117,33 @@ def test_one_filtered_repo_files(allow_patterns: list[str], expected_bool: bool) ) +@pytest.mark.parametrize( + ("cache_result", "should_download"), + [ + # HF Hub recorded a prior 404: don't re-probe the Hub. + (_CACHED_NO_EXIST, False), + # File not in cache and existence unknown: preserve download behavior. + (None, True), + ], +) +def test_get_hf_file_to_dict_honors_no_exist_marker( + cache_result: object, should_download: bool +): + with ( + patch( + "vllm.transformers_utils.repo_utils.try_to_load_from_cache", + MagicMock(return_value=cache_result), + ), + patch( + "vllm.transformers_utils.repo_utils._try_download_from_hf_hub", + MagicMock(return_value=None), + ) as mock_download, + ): + result = get_hf_file_to_dict("processor_config.json", "some/repo") + assert result is None + assert mock_download.call_count == int(should_download) + + @pytest.mark.parametrize( ("files", "expected_bool"), [ diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 2d8a32ef3d5..d6366407247 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -896,9 +896,9 @@ def get_sentence_transformer_tokenizer_config( encoder_dict = None for config_file in sentence_transformer_config_files: - if ( - try_get_local_file(model=model, file_name=config_file, revision=revision) - is not None + if isinstance( + try_get_local_file(model=model, file_name=config_file, revision=revision), + Path, ): encoder_dict = get_hf_file_to_dict(config_file, model, revision) if encoder_dict: diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index 8385057e911..5506af4cac8 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -9,7 +9,7 @@ import time from collections.abc import Callable from functools import cache from pathlib import Path -from typing import TypeVar +from typing import Any, TypeVar import huggingface_hub from huggingface_hub import HfApi, try_to_load_from_cache @@ -218,8 +218,11 @@ def file_or_path_exists( # NB: file_exists will only check for the existence of the config file on # hf_hub. This will fail in offline mode. - # Call HF to check if the file exists - return file_exists(str(model), config_name, revision=revision) + if cached_filepath is None: + # The config file is not cached - check if it exists on hf_hub + return file_exists(str(model), config_name, revision=revision) + # The config file is known to not exist in cache - we can return False + return False def get_model_path(model: str | Path, revision: str | None = None): @@ -288,7 +291,7 @@ def get_hf_file_bytes( if file_path is None: file_path = _try_download_from_hf_hub(model, file_name, revision) - if file_path is not None and file_path.is_file(): + if isinstance(file_path, Path) and file_path.is_file(): with open(file_path, "rb") as file: return file.read() @@ -297,7 +300,20 @@ def get_hf_file_bytes( def try_get_local_file( model: str | Path, file_name: str, revision: str | None = "main" -) -> Path | None: +) -> Path | Any | None: + """ + Try to get a local file from the HuggingFace repository. + + The possible return values are: + + - A `Path` object if the local file is found + - The `huggingface_hub._CACHED_NO_EXIST` sentinel if the file is known to not exist + - `None` if the file is not found and we cannot determine if it exists or not + + Callers of this method should handle the `_CACHED_NO_EXIST` sentinel appropriately. + Checking if the return value `is not None` is not sufficient because it does not + distinguish between the file not existing and the file not being found. + """ file_path = Path(model) / file_name if file_path.is_file(): return file_path @@ -308,6 +324,7 @@ def try_get_local_file( ) if isinstance(cached_filepath, str): return Path(cached_filepath) + return cached_filepath except ValueError: ... return None @@ -335,7 +352,7 @@ def get_hf_file_to_dict( if file_path is None: file_path = _try_download_from_hf_hub(model, file_name, revision) - if file_path is not None and file_path.is_file(): + if isinstance(file_path, Path) and file_path.is_file(): with open(file_path) as file: return json.load(file) From af16446bf39de047ab57649c933063cf1cbf1e50 Mon Sep 17 00:00:00 2001 From: Brandon Pelfrey Date: Fri, 26 Jun 2026 17:32:51 -0700 Subject: [PATCH 058/138] Vram semaphore infra (#44465) Signed-off-by: Brandon Pelfrey Co-authored-by: Roger Wang --- requirements/cuda.txt | 1 + tests/multimodal/test_gpu_ipc_memory.py | 145 ++++++++++ tests/multimodal/test_video.py | 234 ++++++++++++++++ tests/v1/worker/test_gpu_worker.py | 116 ++++++++ vllm/config/model.py | 4 + vllm/config/multimodal.py | 10 + vllm/engine/arg_utils.py | 6 + vllm/multimodal/gpu_ipc_memory.py | 147 ++++++++++ vllm/multimodal/video.py | 344 +++++++++++++++++++++++- vllm/renderers/base.py | 11 + vllm/v1/worker/gpu_worker.py | 84 +++++- 11 files changed, 1087 insertions(+), 15 deletions(-) create mode 100644 tests/multimodal/test_gpu_ipc_memory.py create mode 100644 tests/v1/worker/test_gpu_worker.py create mode 100644 vllm/multimodal/gpu_ipc_memory.py diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 19a1f63dd91..124dae4846d 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -8,6 +8,7 @@ torch==2.11.0 torchaudio==2.11.0 # These must be updated alongside torch torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version +PyNvVideoCodec==2.0.4 # FlashInfer should be updated together with the Dockerfile flashinfer-python==0.6.12 flashinfer-cubin==0.6.12 diff --git a/tests/multimodal/test_gpu_ipc_memory.py b/tests/multimodal/test_gpu_ipc_memory.py new file mode 100644 index 00000000000..bc6bf031fec --- /dev/null +++ b/tests/multimodal/test_gpu_ipc_memory.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +import time + +import pytest + +from vllm.multimodal.gpu_ipc_memory import ( + MultiModalGPUMemoryPool, + get_mm_gpu_ipc_pool, + maybe_init_mm_gpu_ipc_pool, + set_mm_gpu_ipc_pool, +) +from vllm.utils.mem_constants import GiB_bytes + + +def test_acquire_release_accounting(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + assert pool.available_bytes == 100 + + lease = pool.acquire(40) + assert pool.available_bytes == 60 + + lease.release() + assert pool.available_bytes == 100 + + +def test_acquire_too_large_raises(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + with pytest.raises(ValueError): + pool.acquire(101) + # Nothing should have been reserved. + assert pool.available_bytes == 100 + + +def test_negative_acquire_raises(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + with pytest.raises(ValueError): + pool.acquire(-1) + + +def test_double_release_is_noop(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + lease = pool.acquire(50) + lease.release() + assert pool.available_bytes == 100 + # Releasing again must not inflate the pool past its capacity. + lease.release() + assert pool.available_bytes == 100 + + +def test_context_manager_releases_on_exception(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + with pytest.raises(RuntimeError), pool.acquire(50): + assert pool.available_bytes == 50 + raise RuntimeError("boom") + assert pool.available_bytes == 100 + + +def test_acquire_blocks_until_release(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + first = pool.acquire(80) + + acquired = threading.Event() + + def waiter(): + # Needs 50 bytes but only 20 are free; must block until `first` + # is released. + with pool.acquire(50): + acquired.set() + + t = threading.Thread(target=waiter) + t.start() + + # The waiter cannot proceed yet. + assert not acquired.wait(timeout=0.2) + + # Releasing the first lease frees enough budget to unblock the waiter. + first.release() + assert acquired.wait(timeout=2.0) + t.join(timeout=2.0) + assert not t.is_alive() + assert pool.available_bytes == 100 + + +def test_concurrent_acquires_serialize(): + pool = MultiModalGPUMemoryPool(total_bytes=100) + # Each task needs 60 bytes, so only one can hold the budget at a time. + in_section = [] + max_concurrent = 0 + lock = threading.Lock() + + def task(): + nonlocal max_concurrent + with pool.acquire(60): + with lock: + in_section.append(1) + max_concurrent = max(max_concurrent, len(in_section)) + time.sleep(0.05) + with lock: + in_section.pop() + + threads = [threading.Thread(target=task) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5.0) + assert not t.is_alive() + + assert max_concurrent == 1 + assert pool.available_bytes == 100 + + +def test_zero_total_bytes_rejected(): + with pytest.raises(ValueError): + MultiModalGPUMemoryPool(total_bytes=0) + + +def test_global_pool_accessor(): + try: + assert maybe_init_mm_gpu_ipc_pool(0) is None + assert get_mm_gpu_ipc_pool() is None + + pool = maybe_init_mm_gpu_ipc_pool(2) + assert pool is not None + assert get_mm_gpu_ipc_pool() is pool + assert pool.total_bytes == 2 * GiB_bytes + finally: + set_mm_gpu_ipc_pool(None) + + +def test_global_pool_splits_budget_across_api_processes(): + try: + pool = maybe_init_mm_gpu_ipc_pool(2, api_process_count=4) + assert pool is not None + assert get_mm_gpu_ipc_pool() is pool + assert pool.total_bytes == GiB_bytes // 2 + finally: + set_mm_gpu_ipc_pool(None) + + +def test_global_pool_rejects_invalid_api_process_count(): + with pytest.raises(ValueError): + maybe_init_mm_gpu_ipc_pool(2, api_process_count=0) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 694eb392c48..6fccc926a21 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import itertools +import sys +import threading +from contextlib import ExitStack, contextmanager from pathlib import Path import numpy as np @@ -11,10 +15,15 @@ from transformers.video_utils import VideoMetadata from vllm.assets.base import get_vllm_public_assets from vllm.multimodal.video import ( + PYNVVIDEOCODEC_DECODER_CACHE_SIZE, + PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, + PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, DynamicVideoBackend, GLM46VVideoBackend, Molmo2VideoBackend, + PyNvVideoCodecDecoderSlot, + PyNvVideoCodecVideoBackend, Qwen2VLVideoBackend, Qwen3VLVideoBackend, VideoLoader, @@ -22,6 +31,7 @@ from vllm.multimodal.video import ( VideoTargetMetadata, get_video_loader_backend_for_processor, ) +from vllm.platforms import current_platform from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config from .utils import create_long_gop_video, create_video_from_image @@ -65,6 +75,230 @@ def test_video_loader_type_doesnt_exist(): VIDEO_LOADER_REGISTRY.load("non_existing_video_loader") +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_pynvvideocodec_backend_accounts_raw_decoded_frames( + monkeypatch: pytest.MonkeyPatch, +): + decoder_cache_sizes = [] + + class FakeMetadata: + width = 10 + height = 20 + average_fps = 5.0 + duration = 2.0 + + class FakeDecoder: + def __init__(self, *args, **kwargs): + decoder_cache_sizes.append(kwargs["decoder_cache_size"]) + + def __len__(self): + return 10 + + def get_stream_metadata(self): + return FakeMetadata() + + class FakeNvc: + class OutputColorType: + RGB = "rgb" + + SimpleDecoder = FakeDecoder + + class RecordingPool: + def __init__(self): + self.acquired: list[int] = [] + + @contextmanager + def acquire(self, size: int): + self.acquired.append(size) + yield + + def fake_decode(cls, file_path: str, frame_idx: list[int], nvc): + return np.zeros((len(frame_idx), 20, 10, 3), dtype=np.uint8) + + pool = RecordingPool() + monkeypatch.setitem(sys.modules, "PyNvVideoCodec", FakeNvc) + monkeypatch.setattr( + "vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool + ) + monkeypatch.setattr( + PyNvVideoCodecVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode) + ) + + loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND) + frames, metadata = loader.load_bytes(b"fake video", num_frames=4) + + assert frames.shape == (4, 20, 10, 3) + assert pool.acquired == [4 * 20 * 10 * 3] + assert decoder_cache_sizes == [PYNVVIDEOCODEC_DECODER_CACHE_SIZE] + assert metadata["video_backend"] == PYNVVIDEOCODEC_VIDEO_BACKEND + assert metadata["frames_indices"] == [0, 3, 6, 9] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_pynvvideocodec_codec_uses_dynamic_sampling_strategy( + monkeypatch: pytest.MonkeyPatch, +): + decoded_indices = [] + + class FakeMetadata: + width = 10 + height = 20 + average_fps = 5.0 + duration = 2.0 + + class FakeDecoder: + def __init__(self, *args, **kwargs): + pass + + def __len__(self): + return 10 + + def get_stream_metadata(self): + return FakeMetadata() + + class FakeNvc: + class OutputColorType: + RGB = "rgb" + + SimpleDecoder = FakeDecoder + + class RecordingPool: + def __init__(self): + self.acquired: list[int] = [] + + @contextmanager + def acquire(self, size: int): + self.acquired.append(size) + yield + + def fake_decode(cls, file_path: str, frame_idx: list[int], nvc): + decoded_indices.append(frame_idx) + return np.zeros((len(frame_idx), 20, 10, 3), dtype=np.uint8) + + pool = RecordingPool() + monkeypatch.setitem(sys.modules, "PyNvVideoCodec", FakeNvc) + monkeypatch.setattr( + "vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool + ) + monkeypatch.setattr( + DynamicVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode) + ) + + loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic") + frames, metadata = loader.load_bytes( + b"fake video", + fps=2, + max_duration=1, + backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + + assert frames.shape == (2, 20, 10, 3) + assert decoded_indices == [[0, 9]] + assert pool.acquired == [2 * 20 * 10 * 3] + assert metadata["video_backend"] == f"{PYNVVIDEOCODEC_VIDEO_BACKEND}_dynamic" + assert metadata["frames_indices"] == [0, 9] + + +def test_pynvvideocodec_decoder_slots_are_bounded(monkeypatch: pytest.MonkeyPatch): + class FakeSlot: + pass + + create_count = 0 + old_slots = PyNvVideoCodecVideoBackend._decoder_slots + old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots + old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond + try: + PyNvVideoCodecVideoBackend._decoder_slots = [] + PyNvVideoCodecVideoBackend._active_decoder_slots = 0 + PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() + + def fake_create_slot(cls): + nonlocal create_count + create_count += 1 + return FakeSlot() + + monkeypatch.setattr( + PyNvVideoCodecVideoBackend, + "_create_decoder_slot", + classmethod(fake_create_slot), + ) + + borrowed = threading.Event() + seen_slots = [] + + with ExitStack() as stack: + retained_slots = [ + stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot()) + for _ in range(PYNVVIDEOCODEC_MAX_RETAINED_DECODERS) + ] + + def borrow_extra_slot(): + with PyNvVideoCodecVideoBackend._borrow_decoder_slot() as extra_slot: + seen_slots.append(extra_slot) + borrowed.set() + + thread = threading.Thread(target=borrow_extra_slot) + thread.start() + assert not borrowed.wait(timeout=0.2) + + assert borrowed.wait(timeout=2.0) + thread.join(timeout=2.0) + assert not thread.is_alive() + + assert seen_slots[0] in retained_slots + assert create_count == PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + finally: + PyNvVideoCodecVideoBackend._decoder_slots = old_slots + PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots + PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond + + +def test_pynvvideocodec_decoder_slot_retains_simple_decoder(): + events: list[tuple[object, ...]] = [] + + class FakeStream: + cuda_stream = "cuda-stream" + + class FakeDecoder: + def __init__(self, file_path: str, **kwargs): + events.append( + ( + "create", + file_path, + kwargs["gpu_id"], + kwargs["cuda_stream"], + kwargs["decoder_cache_size"], + ) + ) + + def reconfigure_decoder(self, file_path: str): + events.append(("reconfigure", file_path)) + + class FakeNvc: + class OutputColorType: + RGB = "rgb" + + SimpleDecoder = FakeDecoder + + slot = PyNvVideoCodecDecoderSlot(FakeStream()) + + decoder = slot.get_decoder("first.mp4", FakeNvc, device_index=7) + assert slot.get_decoder("first.mp4", FakeNvc, device_index=7) is decoder + assert slot.get_decoder("second.mp4", FakeNvc, device_index=7) is decoder + + assert events == [ + ( + "create", + "first.mp4", + 7, + "cuda-stream", + PYNVVIDEOCODEC_DECODER_CACHE_SIZE, + ), + ("reconfigure", "second.mp4"), + ] + assert slot.source_path == "second.mp4" + + # ============================================================================ # Video Processor โ†’ Video Loader Tests (via model repo) # ============================================================================ diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py new file mode 100644 index 00000000000..31be4a8402f --- /dev/null +++ b/tests/v1/worker/test_gpu_worker.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +import vllm.v1.worker.gpu_worker as gpu_worker_module +from vllm.multimodal.video import ( + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, + PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, + PYNVVIDEOCODEC_VIDEO_BACKEND, +) +from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.worker.gpu_worker import Worker + + +def _worker_with_mm_config( + mm_config: SimpleNamespace, + *, + api_process_count: int = 1, +) -> Worker: + worker = object.__new__(Worker) + worker.model_config = SimpleNamespace(multimodal_config=mm_config) + worker.parallel_config = SimpleNamespace(_api_process_count=api_process_count) + return worker + + +def _mm_config( + *, + mm_ipc_gpu_memory_gb: float = 0, + video_backend: str | None = None, +) -> SimpleNamespace: + video_kwargs = {} if video_backend is None else {"video_backend": video_backend} + return SimpleNamespace( + mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, + media_io_kwargs={"video": video_kwargs} if video_kwargs else {}, + ) + + +def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int: + return api_process_count * ( + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES + ) + + +@pytest.mark.parametrize("video_backend", [None, "opencv"]) +def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only( + monkeypatch: pytest.MonkeyPatch, + video_backend: str | None, +): + monkeypatch.setattr( + gpu_worker_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + "opencv", + ) + worker = _worker_with_mm_config( + _mm_config(mm_ipc_gpu_memory_gb=0.25, video_backend=video_backend) + ) + + assert worker._reserve_mm_ipc_gpu_memory(GiB_bytes) == int(0.75 * GiB_bytes) + + +def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + gpu_worker_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + "opencv", + ) + worker = _worker_with_mm_config( + _mm_config( + mm_ipc_gpu_memory_gb=0.25, + video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + ) + available_bytes = 4 * GiB_bytes + + assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( + available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget() + ) + + +def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + gpu_worker_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + worker = _worker_with_mm_config(_mm_config()) + available_bytes = 4 * GiB_bytes + + assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( + available_bytes - _pynvvideocodec_decoder_budget() + ) + + +def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + gpu_worker_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + PYNVVIDEOCODEC_VIDEO_BACKEND, + ) + worker = _worker_with_mm_config(_mm_config(), api_process_count=3) + available_bytes = 8 * GiB_bytes + + assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( + available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3) + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index c7736f985df..fecb26aa7e0 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -351,6 +351,7 @@ class ModelConfig: skip_mm_profiling: InitVar[bool | None] = None video_pruning_rate: InitVar[float | None] = None mm_tensor_ipc: InitVar[MMTensorIPC] = None + mm_ipc_gpu_memory_gb: InitVar[float | None] = None def compute_hash(self) -> str: """ @@ -397,6 +398,7 @@ class ModelConfig: "mm_encoder_tp_mode", "interleave_mm_strings", "skip_mm_profiling", + "mm_ipc_gpu_memory_gb", } from vllm.config.utils import get_hash_factors, hash_factors @@ -477,6 +479,7 @@ class ModelConfig: skip_mm_profiling: bool | None, video_pruning_rate: float | None, mm_tensor_ipc: MMTensorIPC, + mm_ipc_gpu_memory_gb: float | None, ) -> None: # Keep set served_model_name before maybe_model_redirect(self.model) self.served_model_name = get_served_model_name( @@ -690,6 +693,7 @@ class ModelConfig: skip_mm_profiling=skip_mm_profiling, video_pruning_rate=video_pruning_rate, mm_tensor_ipc=mm_tensor_ipc, + mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, ) mm_config_kwargs = { diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 56333b1116c..150d58cf1f7 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -197,6 +197,16 @@ class MultiModalConfig: - "direct_rpc": Use msgspec serialization via RPC - "torch_shm": Use torch.multiprocessing shared memory for zero-copy IPC Defaults to "direct_rpc". """ + mm_ipc_gpu_memory_gb: float = Field(default=0, ge=0) + """Amount of GPU memory (in GiB) sequestered on the engine's device for + GPU-side multimodal work in the API-server (frontend) process, such as + hardware video decoding. + + This budget is carved out of the engine's KV-cache memory so the headroom + physically exists, and frontend GPU decode paths acquire from a blocking + byte-counting semaphore of this size before allocating on the device. + + Set to `0` (default) to disable frontend GPU multimodal memory gating.""" @field_validator("limit_per_prompt", mode="before") @classmethod diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 8cc219264f3..efdd7696fdc 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -582,6 +582,7 @@ class EngineArgs: skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc + mm_ipc_gpu_memory_gb: float = MultiModalConfig.mm_ipc_gpu_memory_gb # LoRA fields enable_lora: bool = False max_loras: int = LoRAConfig.max_loras @@ -1294,6 +1295,10 @@ class EngineArgs: multimodal_group.add_argument( "--mm-tensor-ipc", **multimodal_kwargs["mm_tensor_ipc"] ) + multimodal_group.add_argument( + "--mm-ipc-gpu-memory-gb", + **multimodal_kwargs["mm_ipc_gpu_memory_gb"], + ) # LoRA related configs lora_kwargs = get_kwargs(LoRAConfig) @@ -1660,6 +1665,7 @@ class EngineArgs: logits_processors=self.logits_processors, video_pruning_rate=self.video_pruning_rate, mm_tensor_ipc=self.mm_tensor_ipc, + mm_ipc_gpu_memory_gb=self.mm_ipc_gpu_memory_gb, io_processor_plugin=self.io_processor_plugin, renderer_num_workers=self.renderer_num_workers, ) diff --git a/vllm/multimodal/gpu_ipc_memory.py b/vllm/multimodal/gpu_ipc_memory.py new file mode 100644 index 00000000000..15b912064a8 --- /dev/null +++ b/vllm/multimodal/gpu_ipc_memory.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Admission control for frontend GPU-side multimodal work. + +When multimodal media is decoded on the GPU in the API-server (frontend) +process, the decoded buffers compete for the same device memory that the +engine reserves for weights, activations, and the KV cache. To keep the +frontend's GPU usage within a sequestered budget (see +``MultiModalConfig.mm_ipc_gpu_memory_gb``), decode paths acquire the number of +bytes they need from a process-global :class:`MultiModalGPUMemoryPool` before +allocating on the device and release them once the device memory is freed. + +The pool is a simple byte-counting semaphore: ``acquire`` blocks until enough +budget is free, so concurrent requests serialize rather than oversubscribe the +GPU. It lives only in the frontend process; the engine carves the matching +amount out of its KV-cache budget so the headroom physically exists. +""" + +import threading + +from vllm.logger import init_logger +from vllm.utils.mem_constants import GiB_bytes + +logger = init_logger(__name__) + + +class MultiModalGPUMemoryLease: + """A handle for bytes acquired from a :class:`MultiModalGPUMemoryPool`. + + Releasing is idempotent and the lease doubles as a context manager so the + budget is returned even if the decode raises. + """ + + def __init__(self, pool: "MultiModalGPUMemoryPool", lease_id: int, nbytes: int): + self.lease_id = lease_id + self.nbytes = nbytes + self._pool = pool + + def release(self) -> None: + self._pool._release(self) + + def __enter__(self) -> "MultiModalGPUMemoryLease": + return self + + def __exit__(self, *exc_info) -> None: + self.release() + + +class MultiModalGPUMemoryPool: + """Blocking byte-counting semaphore for frontend GPU multimodal memory. + + Thread-safe in both directions: ``acquire`` (blocking) and ``release`` are + typically called from the renderer's multimodal executor threads. + """ + + def __init__(self, total_bytes: int): + if total_bytes <= 0: + raise ValueError(f"total_bytes must be positive, got {total_bytes}") + self._total_bytes = total_bytes + self._available = total_bytes + self._cond = threading.Condition() + self._next_lease_id = 0 + # Outstanding lease ids, so a double release is a no-op. + self._outstanding: set[int] = set() + + @property + def total_bytes(self) -> int: + return self._total_bytes + + @property + def available_bytes(self) -> int: + with self._cond: + return self._available + + def acquire(self, nbytes: int) -> MultiModalGPUMemoryLease: + """Reserve ``nbytes``, blocking until that much budget is free. + + Raises ``ValueError`` if ``nbytes`` exceeds the pool's total capacity, + since such a request could never be satisfied. + """ + if nbytes < 0: + raise ValueError(f"Cannot acquire negative bytes: {nbytes}") + if nbytes > self._total_bytes: + raise ValueError( + f"Multimodal GPU decode requested {nbytes} bytes, which exceeds " + f"the total pool size of {self._total_bytes} bytes. Increase " + f"--mm-ipc-gpu-memory-gb or reduce the multimodal input size." + ) + with self._cond: + while self._available < nbytes: + self._cond.wait() + self._available -= nbytes + lease_id = self._next_lease_id + self._next_lease_id += 1 + self._outstanding.add(lease_id) + return MultiModalGPUMemoryLease(self, lease_id, nbytes) + + def _release(self, lease: MultiModalGPUMemoryLease) -> None: + with self._cond: + if lease.lease_id not in self._outstanding: + # Already released โ€” idempotent. + return + self._outstanding.discard(lease.lease_id) + self._available += lease.nbytes + self._cond.notify_all() + + +_GLOBAL_POOL: MultiModalGPUMemoryPool | None = None + + +def set_mm_gpu_ipc_pool(pool: MultiModalGPUMemoryPool | None) -> None: + """Install the process-global pool (frontend process only).""" + global _GLOBAL_POOL + _GLOBAL_POOL = pool + + +def get_mm_gpu_ipc_pool() -> MultiModalGPUMemoryPool | None: + """Return the process-global pool, or ``None`` when gating is disabled.""" + return _GLOBAL_POOL + + +def maybe_init_mm_gpu_ipc_pool( + mm_ipc_gpu_memory_gb: float, + api_process_count: int = 1, +) -> MultiModalGPUMemoryPool | None: + """Create and install the global pool from the configured GiB budget. + + Returns ``None`` (and leaves gating disabled) when the budget is 0. When + multiple API-server processes share one engine, each process gets an equal + slice of the user-provided frontend budget. + """ + if mm_ipc_gpu_memory_gb <= 0: + set_mm_gpu_ipc_pool(None) + return None + if api_process_count <= 0: + raise ValueError(f"api_process_count must be positive, got {api_process_count}") + total_bytes = int(mm_ipc_gpu_memory_gb * GiB_bytes) // api_process_count + pool = MultiModalGPUMemoryPool(total_bytes) + set_mm_gpu_ipc_pool(pool) + logger.info( + "Initialized multimodal GPU IPC memory pool with %d bytes for this API " + "process (%.2f GiB total budget across %d API process(es)).", + total_bytes, + mm_ipc_gpu_memory_gb, + api_process_count, + ) + return pool diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 700bd0802c5..8cd4870026e 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +import os +import tempfile +import threading from abc import abstractmethod +from contextlib import contextmanager, suppress from io import BytesIO from typing import Any, ClassVar, Literal, NamedTuple, cast @@ -11,6 +15,7 @@ import torch from vllm.logger import init_logger from vllm.utils.import_utils import PlaceholderModule +from vllm.utils.mem_constants import MiB_bytes from vllm.utils.registry import ExtensionManager try: @@ -130,6 +135,14 @@ class VideoSourceMetadata(NamedTuple): duration: float +class PyNvVideoCodecSourceMetadata(NamedTuple): + """Metadata needed before GPU video decode.""" + + source: VideoSourceMetadata + width: int + height: int + + class VideoLoader: @classmethod def compute_frames_index_to_sample( @@ -170,6 +183,57 @@ class VideoLoader: VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() +PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec" +# Fixed upper bound reserved for persistent PyNvVideoCodec decoder surfaces. +PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes +PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2 +PYNVVIDEOCODEC_MAX_RETAINED_DECODERS = 1 +# Per-API-server CUDA context and driver allocation, measured with +# PyNvVideoCodec 2.0.4 on H100. +PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes) + + +class PyNvVideoCodecDecoderSlot: + """A retained PyNv decoder slot and its CUDA stream. + + The decoder is reused across requests: ``reconfigure_decoder`` repoints the + existing decoder at each new source instead of paying a fresh + ``SimpleDecoder`` construction per request. Construction (CUVID parser + + decoder + surface-pool allocation) is the dominant per-request cost, so + reconfiguring is far cheaper. A single decoder serves both metadata + (``len``/``get_stream_metadata``) and frame decode -- no separate + metadata decoder. + """ + + def __init__(self, stream) -> None: + self.stream = stream + self.decoder = None + self.source_path: str | None = None + + def _construct(self, file_path: str, nvc, device_index: int) -> None: + self.decoder = nvc.SimpleDecoder( + file_path, + output_color_type=nvc.OutputColorType.RGB, + use_device_memory=True, + need_scanned_stream_metadata=True, + gpu_id=device_index, + cuda_stream=self.stream.cuda_stream, + decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE, + ) + self.source_path = file_path + + def get_decoder(self, file_path: str, nvc, device_index: int): + if self.decoder is None: + self._construct(file_path, nvc, device_index) + elif self.source_path != file_path: + try: + self.decoder.reconfigure_decoder(file_path) + self.source_path = file_path + except Exception: + # reconfigure unsupported/unsafe for this source -> rebuild. + self._construct(file_path, nvc, device_index) + return self.decoder + class OpenCVVideoBackendMixin: @staticmethod @@ -485,15 +549,223 @@ class PyAVVideoBackendMixin: return np.stack(frames_list), valid_indices +class PyNvVideoCodecVideoBackendMixin: + """PyNvVideoCodec utilities for GPU-backed frame decode.""" + + _decoder_slots: ClassVar[list[PyNvVideoCodecDecoderSlot]] = [] + _active_decoder_slots: ClassVar[int] = 0 + _decoder_slot_cond: ClassVar[threading.Condition] = threading.Condition() + _DEVICE_INDEX: ClassVar[int] = 0 + + @classmethod + @abstractmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + raise NotImplementedError + + @classmethod + @abstractmethod + def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: + raise NotImplementedError + + @classmethod + def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot: + import torch + + return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX)) + + @staticmethod + @contextmanager + def _torch_stream_context(stream): + import torch + + torch.accelerator.set_device_index(stream.device.index) + previous_stream = torch.accelerator.current_stream() + torch.accelerator.set_stream(stream) + try: + yield + finally: + torch.accelerator.set_stream(previous_stream) + + @classmethod + @contextmanager + def _borrow_decoder_slot(cls): + create_slot = False + with cls._decoder_slot_cond: + while True: + if cls._decoder_slots: + slot = cls._decoder_slots.pop() + break + if cls._active_decoder_slots < PYNVVIDEOCODEC_MAX_RETAINED_DECODERS: + cls._active_decoder_slots += 1 + create_slot = True + break + cls._decoder_slot_cond.wait() + + if create_slot: + try: + slot = cls._create_decoder_slot() + except Exception: + with cls._decoder_slot_cond: + cls._active_decoder_slots -= 1 + cls._decoder_slot_cond.notify() + raise + + try: + yield slot + finally: + with cls._decoder_slot_cond: + cls._decoder_slots.append(slot) + cls._decoder_slot_cond.notify() + + @staticmethod + def _metadata_value(metadata, *names: str, default=None): + for name in names: + value = getattr(metadata, name, None) + if value is not None: + return value + return default + + @classmethod + def _read_source_metadata( + cls, + file_path: str, + nvc, + ) -> PyNvVideoCodecSourceMetadata: + with cls._borrow_decoder_slot() as decoder_slot: + with cls._torch_stream_context(decoder_slot.stream): + decoder = decoder_slot.get_decoder( + file_path, nvc, device_index=cls._DEVICE_INDEX + ) + metadata = decoder.get_stream_metadata() + total_frames_num = len(decoder) + width = int(cls._metadata_value(metadata, "width", default=0)) + height = int(cls._metadata_value(metadata, "height", default=0)) + original_fps = float( + cls._metadata_value( + metadata, + "average_fps", + "avg_frame_rate", + "frame_rate", + "frameRate", + default=0.0, + ) + ) + duration = float( + cls._metadata_value(metadata, "duration", default=0.0) + or (total_frames_num / original_fps if original_fps > 0 else 0.0) + ) + if total_frames_num <= 0: + raise ValueError("Could not determine video frame count") + if width <= 0 or height <= 0: + raise ValueError("Could not determine video dimensions") + return PyNvVideoCodecSourceMetadata( + source=VideoSourceMetadata(total_frames_num, original_fps, duration), + width=width, + height=height, + ) + + @classmethod + def _decode_to_pinned_host( + cls, + file_path: str, + frame_idx: list[int], + nvc, + ) -> npt.NDArray: + import torch + + if not frame_idx: + return np.empty((0,), dtype=np.uint8) + + with cls._borrow_decoder_slot() as decoder_slot: + stream = decoder_slot.stream + with cls._torch_stream_context(stream): + decoder = decoder_slot.get_decoder( + file_path, nvc, device_index=cls._DEVICE_INDEX + ) + decoded_frames = decoder.get_batch_frames_by_index(frame_idx) + if len(decoded_frames) < len(frame_idx): + logger.warning( + "pynvvideocodec video loading: expected %d frames but got %d.", + len(frame_idx), + len(decoded_frames), + ) + torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames] + if not torch_frames: + return np.empty((0,), dtype=np.uint8) + device_frames = torch.stack(torch_frames) + if device_frames.ndim != 4: + raise ValueError( + "PyNvVideoCodec returned frames with unexpected shape " + f"{tuple(device_frames.shape)}" + ) + device_frames = device_frames.permute(0, 3, 1, 2).contiguous() + host_frames = torch.empty( + device_frames.shape, + dtype=device_frames.dtype, + device="cpu", + pin_memory=True, + ) + host_frames.copy_(device_frames, non_blocking=True) + stream.synchronize() + host_array = host_frames.numpy() + del decoded_frames, torch_frames, device_frames + return host_array + + @classmethod + def decode_frames_pynvvideocodec( + cls, + data: bytes, + target: VideoTargetMetadata, + **kwargs, + ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + import PyNvVideoCodec as nvc + + from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool + + temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4") + try: + with os.fdopen(temp_fd, "wb") as temp_file: + temp_file.write(data) + + gpu_source = cls._read_source_metadata(temp_path, nvc) + source = cls._prepare_source(gpu_source.source) + frame_idx = cls.compute_frames_index_to_sample( + source=source, target=target, **kwargs + ) + raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3 + pool = get_mm_gpu_ipc_pool() + if pool is None or raw_frame_bytes == 0: + frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc) + else: + with pool.acquire(raw_frame_bytes): + frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc) + finally: + with suppress(FileNotFoundError): + os.unlink(temp_path) + + valid_frame_indices = frame_idx[: int(frames.shape[0])] + return frames, source, frame_idx, valid_frame_indices + + @VIDEO_LOADER_REGISTRY.register("opencv") -class VideoBackend(VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin): +class VideoBackend( + VideoLoader, + OpenCVVideoBackendMixin, + PyAVVideoBackendMixin, + PyNvVideoCodecVideoBackendMixin, +): """Uniform-sampling video backend. Samples ``num_frames`` uniformly across the video (or one frame every ``1/fps`` seconds, whichever produces fewer frames). The decoding codec - is selected via the ``backend`` kwarg (``"opencv"`` or ``"pyav"``), - which can be passed through ``--media-io-kwargs``. Defaults to - ``"pyav"`` for concurrent decoding. + is selected via the ``backend`` kwarg (``"opencv"``, ``"pyav"``, or + ``"pynvvideocodec"``), which can be passed through + ``--media-io-kwargs``. Defaults to ``"opencv"``. """ _sampling_suffix: ClassVar[str] = "" @@ -538,7 +810,7 @@ class VideoBackend(VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. @@ -551,7 +823,8 @@ class VideoBackend(VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin): dynamic subclass; ignored here. frame_recovery: Enable forward-scan recovery for failed frames. Only honored by the OpenCV codec. - backend: Decoding codec โ€” ``"opencv"`` or ``"pyav"`` . + backend: Decoding codec โ€” ``"opencv"``, ``"pyav"``, or + ``"pynvvideocodec"``. Returns: Tuple of ``(frames_array, metadata_dict)``. @@ -584,10 +857,21 @@ class VideoBackend(VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin): frames, valid = cls.decode_frames( container, frame_idx, source.original_fps, source.duration ) + elif backend == PYNVVIDEOCODEC_VIDEO_BACKEND: + if frame_recovery: + raise ValueError( + "frame_recovery is not supported for " + f"`{PYNVVIDEOCODEC_VIDEO_BACKEND}` backend" + ) + frames, source, frame_idx, valid = cls.decode_frames_pynvvideocodec( + data, + target, + **kwargs, + ) else: raise ValueError( f"Unknown video codec backend {backend!r}; " - "valid options: 'opencv', 'pyav'." + "valid options: 'opencv', 'pyav', 'pynvvideocodec'." ) if len(valid) < len(frame_idx): @@ -605,6 +889,40 @@ class VideoBackend(VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin): ) +@VIDEO_LOADER_REGISTRY.register(PYNVVIDEOCODEC_VIDEO_BACKEND) +class PyNvVideoCodecVideoBackend(VideoBackend): + """Hardware-accelerated video backend using PyNvVideoCodec. + + The backend first opens the stream only to read metadata and compute the + sampled frame indices. It then acquires the raw decoded RGB byte count from + the process-local multimodal GPU memory pool before decoding the selected + frames into VRAM. Decoded frames are copied into pinned host memory before + the lease is released, so downstream preprocessing continues to receive a + CPU ``np.ndarray`` in NHWC RGB format. + """ + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = -1, + max_duration: int = 300, + frame_recovery: bool = False, + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + kwargs.pop("backend", None) + return super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + **kwargs, + ) + + @VIDEO_LOADER_REGISTRY.register( "qwen3_vl", video_processor="Qwen3VLVideoProcessor", @@ -640,7 +958,7 @@ class Qwen3VLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -719,7 +1037,7 @@ class Qwen2VLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -811,7 +1129,7 @@ class DynamicVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -936,7 +1254,7 @@ class GLM46VVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1034,7 +1352,7 @@ class GLMGAVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( @@ -1353,7 +1671,7 @@ class NemotronVLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav"] = "opencv", + backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 00cbec33d6f..a0f2508ccc7 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -29,6 +29,7 @@ from vllm.inputs import ( from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY as mm_registry from vllm.multimodal.cache import BaseMultiModalProcessorCache +from vllm.multimodal.gpu_ipc_memory import maybe_init_mm_gpu_ipc_pool from vllm.multimodal.parse import ( MultiModalDataItems, MultiModalUUIDItems, @@ -111,6 +112,16 @@ class BaseRenderer(ABC, Generic[_T]): safe_load_prompt_embeds, executor=self._executor ) if mm_registry.supports_multimodal_inputs(config.model_config): + # Install the process-global GPU memory pool used to gate + # frontend GPU-side multimodal decoding (no-op when the budget + # is 0). Lives in the API-server process only. + mm_config = config.model_config.multimodal_config + if mm_config is not None: + maybe_init_mm_gpu_ipc_pool( + mm_config.mm_ipc_gpu_memory_gb, + config.parallel_config._api_process_count, + ) + mm_processor_cache = mm_registry.processor_cache_from_config(config) with set_default_torch_num_threads(): diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 589a16576eb..9afc2352528 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -50,6 +50,12 @@ from vllm.distributed.weight_transfer import ( from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.model_executor.warmup.kernel_warmup import kernel_warmup +from vllm.multimodal.video import ( + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, + PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, + PYNVVIDEOCODEC_VIDEO_BACKEND, +) from vllm.platforms import current_platform from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper from vllm.sequence import IntermediateTensors @@ -428,7 +434,7 @@ class Worker(WorkerBase): "correspondingly." ) logger.info(msg) - return kv_cache_memory_bytes + return self._reserve_mm_ipc_gpu_memory(kv_cache_memory_bytes) # Execute a forward pass with dummy inputs to profile the memory usage # of the model. @@ -550,7 +556,81 @@ class Worker(WorkerBase): suggested_util, ) - return int(self.available_kv_cache_memory_bytes) + return self._reserve_mm_ipc_gpu_memory( + int(self.available_kv_cache_memory_bytes) + ) + + @staticmethod + def _uses_pynvvideocodec_video_backend(mm_config) -> bool: + video_kwargs = mm_config.media_io_kwargs.get("video", {}) + video_loader_backend = ( + video_kwargs.get("video_backend") or envs.VLLM_VIDEO_LOADER_BACKEND + ) + codec_backend = video_kwargs.get("backend") + return ( + video_loader_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + or codec_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + ) + + def _reserve_mm_ipc_gpu_memory(self, available_kv_cache_memory_bytes: int) -> int: + """Carve frontend multimodal GPU memory out of the KV cache. + + The frontend (API-server) process allocates GPU memory for hardware + multimodal decoding. Raw decoded frames are bounded by + ``mm_ipc_gpu_memory_gb`` and acquired by the frontend semaphore. Some + decoders also keep persistent surfaces around; reserve a fixed upper + bound for those when the corresponding backend is configured. + """ + mm_config = self.model_config.multimodal_config + if mm_config is None: + return available_kv_cache_memory_bytes + + raw_frame_reserved_bytes = int(mm_config.mm_ipc_gpu_memory_gb * GiB_bytes) + # Each api_server_count process runs its OWN decoder surfaces + NVDEC/CUVID + # CUDA context on the GPU, outside this (worker) memory pool. Reserve that + # per-server footprint x api_server_count so gpu_memory_utilization bounds + # TOTAL GPU usage across all API-server processes. Without the multiply, + # HW decode overshoots the budget by ~(api_server_count-1) x per-server and + # OOMs at high gmu, while SW decode (no per-server GPU allocation) does not. + num_api_servers = max(1, getattr(self.parallel_config, "_api_process_count", 1)) + per_server_decoder_bytes = ( + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES + * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES + ) + decoder_reserved_bytes = ( + num_api_servers * per_server_decoder_bytes + if self._uses_pynvvideocodec_video_backend(mm_config) + else 0 + ) + reserved_bytes = raw_frame_reserved_bytes + decoder_reserved_bytes + if reserved_bytes <= 0: + return available_kv_cache_memory_bytes + + remaining = available_kv_cache_memory_bytes - reserved_bytes + if remaining <= 0: + raise ValueError( + f"frontend multimodal GPU decoding reserves " + f"{format_gib(reserved_bytes)} GiB " + f"({format_gib(raw_frame_reserved_bytes)} GiB raw-frame budget, " + f"{format_gib(decoder_reserved_bytes)} GiB decoder cache budget), " + f"but only {format_gib(available_kv_cache_memory_bytes)} GiB is " + "available for the KV cache. Reduce mm_ipc_gpu_memory_gb, use a " + "different video backend, or increase gpu_memory_utilization." + ) + logger.info_once( + "Reserving %s GiB of GPU memory for frontend multimodal decoding " + "(%s GiB raw-frame semaphore budget, %s GiB decoder+CUDA-context " + "across %d API server(s) @ %s GiB/server); " + "KV cache memory reduced to %s GiB.", + format_gib(reserved_bytes), + format_gib(raw_frame_reserved_bytes), + format_gib(decoder_reserved_bytes), + num_api_servers, + format_gib(per_server_decoder_bytes), + format_gib(remaining), + ) + return remaining def get_kv_connector_handshake_metadata( self, From c6dd32a810aa8c4eda5696722c807e53d9f595a5 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 26 Jun 2026 19:42:27 -0700 Subject: [PATCH 059/138] [ModelRunner V2] Support realtime embeddings (#46762) --- tests/v1/worker/test_encoder_runner.py | 6 +-- vllm/model_executor/models/diffusion_gemma.py | 2 +- vllm/v1/worker/gpu/mm/encoder_runner.py | 37 ++++++++++--------- vllm/v1/worker/gpu/model_runner.py | 31 +++++++++------- vllm/v1/worker/gpu/model_states/default.py | 4 ++ vllm/v1/worker/gpu/model_states/interface.py | 8 +++- 6 files changed, 52 insertions(+), 36 deletions(-) diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 79c13a2b96a..70c0426640d 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -53,14 +53,14 @@ def _make_runner( def _gather(runner: EncoderRunner, *, num_scheduled: int, draft_lookahead: int): - # Single prefilling request, computed_prefill=0, prefill_len large. + # Single prefilling request, num_computed_tokens=0, prefill_len large. return runner.gather_mm_embeddings( req_ids=["req0"], total_num_scheduled_tokens=num_scheduled, num_scheduled_tokens=np.array([num_scheduled]), query_start_loc=np.array([0]), prefill_lens=np.array([1000]), - computed_prefill_lens=np.array([0]), + num_computed_tokens=np.array([0]), draft_lookahead=draft_lookahead, ) @@ -147,7 +147,7 @@ def test_multi_request_batch_gathers_per_request(draft_lookahead): num_scheduled_tokens=np.array([8, 8]), query_start_loc=np.array([0, 8]), prefill_lens=np.array([1000, 1000]), - computed_prefill_lens=np.array([0, 0]), + num_computed_tokens=np.array([0, 0]), draft_lookahead=draft_lookahead, ) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index e28e2720a7f..6121e55dab8 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -878,7 +878,7 @@ class DiffusionGemmaModelState(ModelState): scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch, req_states: RequestState, - ) -> torch.Tensor: + ) -> torch.Tensor | None: if not self.supports_mm_inputs: return None diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index f0e99fae1f5..48a3af25053 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -3,7 +3,7 @@ import numpy as np import torch -from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime from vllm.multimodal.inputs import MultiModalKwargsItem from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -26,6 +26,7 @@ class EncoderRunner: self.encoder_cache = encoder_cache self.dtype = dtype self.device = device + self.is_realtime = supports_realtime(model) self.inputs_embeds = torch.zeros( max_num_tokens, hidden_size, dtype=dtype, device=device @@ -67,30 +68,32 @@ class EncoderRunner: num_scheduled_tokens: np.ndarray, query_start_loc: np.ndarray, prefill_lens: np.ndarray, - computed_prefill_lens: np.ndarray, + num_computed_tokens: np.ndarray, draft_lookahead: int = 0, ) -> tuple[list[torch.Tensor], torch.Tensor]: if draft_lookahead: - computed_prefill_lens = computed_prefill_lens + draft_lookahead + num_computed_tokens = num_computed_tokens + draft_lookahead - is_prefilling_np = computed_prefill_lens < prefill_lens - if not is_prefilling_np.any(): - # All decode requests, so no need to gather any embeddings. - return [], torch.zeros( - total_num_scheduled_tokens, dtype=torch.bool, device=self.device - ) - - is_prefilling = is_prefilling_np.tolist() - query_start = computed_prefill_lens.tolist() - query_end = (computed_prefill_lens + num_scheduled_tokens).tolist() - - mm_embeds: list[torch.Tensor] = [] is_mm_embed = torch.zeros( total_num_scheduled_tokens, dtype=torch.bool, device="cpu" ) + + # Whether to gather media embeddings this step. + exclude_embeddings: list[bool] | None = None + if not self.is_realtime: + # Non-realtime models only have media embeddings within the prompt. + is_decode = num_computed_tokens >= prefill_lens + if is_decode.all(): + # All decode requests, so no need to gather any embeddings. + return [], is_mm_embed + exclude_embeddings = is_decode.tolist() + + query_start = num_computed_tokens.tolist() + query_end = (num_computed_tokens + num_scheduled_tokens).tolist() + + mm_embeds: list[torch.Tensor] = [] for i, req_id in enumerate(req_ids): - if not is_prefilling[i]: - # OPTIMIZATION: Skip decode requests. + if exclude_embeddings is not None and exclude_embeddings[i]: continue cur_query_start = query_start[i] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index cb46ffc3dc0..927ece4fbac 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1218,20 +1218,25 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.supports_mm_inputs and self.is_first_pp_rank: # Run MM encoder (if needed) and get multimodal embeddings. # Only first PP rank prepares multimodal embeddings. - # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs - # to obtain inputs_embeds, because the compiled model expects this input. - if self.lora_config is not None: - set_active_mm_loras( - model=self.model, - lora_manager=self.lora_manager, - encoder_cache=self.encoder_cache, - req_id_to_index=self.req_states.req_id_to_index, - lora_state=self.lora_state, - scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs, + if dummy_run: + # Obtain mm embeddings of correct shape for compiled model. + inputs_embeds = self.model_state.dummy_inputs_embeds( + input_batch.num_tokens_after_padding + ) + else: + scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs + if self.lora_config is not None: + set_active_mm_loras( + model=self.model, + lora_manager=self.lora_manager, + encoder_cache=self.encoder_cache, + req_id_to_index=self.req_states.req_id_to_index, + lora_state=self.lora_state, + scheduled_encoder_inputs=scheduled_encoder_inputs, + ) + inputs_embeds = self.model_state.get_mm_embeddings( + scheduled_encoder_inputs, input_batch, self.req_states ) - inputs_embeds = self.model_state.get_mm_embeddings( - scheduler_output.scheduled_encoder_inputs, input_batch, self.req_states - ) if inputs_embeds is not None and not self.model.requires_raw_input_tokens: input_ids = None diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 22e6aa00bc9..2e14eb2e7d9 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -57,6 +57,10 @@ class DefaultModelState(ModelState): if self.rope_state is not None: self.rope_state.apply_staged_writes() + def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor: + """Pre-allocated inputs_embeds buffer for dummy runs (contents unused).""" + return self.encoder_runner.inputs_embeds[:num_tokens] + def get_mm_embeddings( self, scheduled_encoder_inputs: dict[str, list[int]], diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index a4c436a423b..c80e19547c0 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -109,17 +109,21 @@ class ModelState(ABC): ) -> torch.Tensor | None: raise NotImplementedError + def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None: + """Pre-allocated inputs_embeds buffer for dummy runs (contents unused).""" + return None + def gather_mm_embeddings( self, input_batch: InputBatch, draft_lookahead: int = 0 ) -> tuple[list[torch.Tensor], torch.Tensor]: - """Gather cached multimodal embeddings for a speculator's draft forward.""" + """Gather cached multimodal embeddings.""" return self.encoder_runner.gather_mm_embeddings( input_batch.req_ids, input_batch.num_tokens, input_batch.num_scheduled_tokens, input_batch.query_start_loc_np, input_batch.prefill_len_np, - input_batch.num_computed_prefill_tokens_np, + input_batch.num_computed_tokens_np, draft_lookahead=draft_lookahead, ) From d0f800811bb8092e6c62a333c05567c3b380ebf8 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 26 Jun 2026 22:42:46 -0400 Subject: [PATCH 060/138] [Build] Update vllm to point to vllm-project/flash-attention commit that builds FA3 with torch stable API. (#46644) --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index ea7ac544b9d..c8b1d689187 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee + GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn From 1a92dfcce4a9433d002631207b5b01e2e95f1077 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C4=81vis?= Date: Sat, 27 Jun 2026 05:43:00 +0300 Subject: [PATCH 061/138] [Build] Show error message when using ROCm with LTO and different compilers (#35232) --- CMakeLists.txt | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 36b8e66f2c6..cbd5583bbfd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -270,6 +270,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP") # set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result -Wno-unused-value") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result -Wno-unused-value") + + # When using LTO then *.cpp files must be compiled with same compiler as used linker + # So if HIP uses clang linker we also must use it + # Otherwise symbols will be missing from .so + if (CMAKE_CXX_FLAGS MATCHES "\-flto") + if(NOT CMAKE_CXX_COMPILER_ID STREQUAL CMAKE_HIP_COMPILER_ID) + message(FATAL_ERROR "LTO is enabled for ROCm build, but the C++ compiler (${CMAKE_CXX_COMPILER_ID}) and HIP compiler (${CMAKE_HIP_COMPILER_ID}) are different which is not supported. " + "Please ensure they are same by setting CXX=${CMAKE_HIP_COMPILER} environment variable. Or alternatively disable LTO.") + endif() + endif() endif() # From 2e058851d39448bf282e64a6aac04466968622e7 Mon Sep 17 00:00:00 2001 From: weizhoublue <45163302+weizhoublue@users.noreply.github.com> Date: Sat, 27 Jun 2026 10:43:17 +0800 Subject: [PATCH 062/138] fix(docker): eliminate race conditions in shared buildkit cache mounts (#44984) --- docker/Dockerfile | 5 +++-- docker/Dockerfile.nightly_torch | 5 +++-- docker/Dockerfile.xpu | 4 ++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index ef166665c5f..c86795586c3 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -311,9 +311,10 @@ ENV CARGO_BUILD_JOBS=4 # Build the release artifacts. Cache cargo registry/git, but not target/, # because stale target metadata can outlive source updates across BuildKit # cache reuse. -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ +RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,target=/root/.cargo/git,sharing=locked \ bash build_rust.sh + #################### RUST BUILD IMAGE #################### #################### CSRC BUILD IMAGE #################### diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 149c265d7e2..0f2ec9f3a2e 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -123,9 +123,10 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ +RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,target=/root/.cargo/git,sharing=locked \ bash build_rust.sh + #################### RUST BUILD IMAGE #################### #################### WHEEL BUILD IMAGE #################### diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 7b87f202769..3bd16e8629b 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -28,8 +28,8 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ +RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,target=/root/.cargo/git,sharing=locked \ bash build_rust.sh FROM ubuntu:24.04 AS vllm-base From 17a71d87020e163a291bd34f91dad3eb05b448e5 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Fri, 26 Jun 2026 22:44:29 -0400 Subject: [PATCH 063/138] [ROCm][CI] Relax fused layernorm quant test tolerances for one-ULP outliers (#46658) Signed-off-by: Divakar Verma --- .../core/test_fused_quant_layernorm.py | 56 +++++++++++++------ tests/kernels/core/test_layernorm.py | 24 +++++--- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 36 ++---------- tests/kernels/utils.py | 29 ++++++++++ 4 files changed, 90 insertions(+), 55 deletions(-) diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py index 07d15e3b1df..255833c48dc 100644 --- a/tests/kernels/core/test_fused_quant_layernorm.py +++ b/tests/kernels/core/test_fused_quant_layernorm.py @@ -8,7 +8,7 @@ import pytest import torch import vllm._custom_ops as ops -from tests.kernels.utils import opcheck +from tests.kernels.utils import fp8_ulp_distance, opcheck from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, @@ -250,30 +250,54 @@ def test_rms_norm( assert ref_out.dtype == quant_dtype assert ops_out.dtype == quant_dtype + + # Per-block bf16 scales: allow a small relative tolerance for a few groups + # whose abs-max flips by one ULP between the fused and reference paths. The + # per-token and fp32 paths stay strict. + relax_block_rocm = ( + group_size is not None + and dtype == torch.bfloat16 + and current_platform.is_rocm() + ) + + def scales_close(rtol: float, atol: float) -> bool: + if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol): + return True + return relax_block_rocm and torch.allclose( + ref_scales, ops_scales, rtol=1e-2, atol=atol + ) + if quant_dtype == torch.int8: - assert torch.allclose(ref_scales, ops_scales, atol=1e-6) + assert scales_close(rtol=1e-5, atol=1e-6) # big atol to account for round-off errors. assert torch.allclose(ref_out, ops_out, atol=1) else: - assert torch.allclose(ref_scales, ops_scales) + assert scales_close(rtol=1e-5, atol=1e-8) a = ref_out.to(dtype=torch.float32) b = ops_out.to(dtype=torch.float32) ok = torch.allclose(a, b, atol=1e-6) if not ok: - # fallback: compare dequantized values with relaxed tolerance - if group_size is None: - a_deq = a * ref_scales.view(-1, 1) - b_deq = b * ops_scales.view(-1, 1) + if relax_block_rocm: + # ULP-flipped group scale can cross an E4M3 tie; tolerate a + # bounded count of isolated fp8 outliers. + ulp = fp8_ulp_distance(ref_out, ops_out) + max_outliers = ulp.numel() // 100_000 + 8 + ok = int((ulp > 0).sum().item()) <= max_outliers else: - a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1) - b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1) - # NOTE: It is possible that some future test cases trigger this - # max diff due to precision issues. If such an error is - # encountered, it's recommended to inspect the differences between - # all corresponding elements from each tensor (e.g. by looping over - # them) and checking how many the max diff error shows up on (just - # a few bad elements should still be considered acceptable). - ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2) + # CUDA (& non-bf16): compare dequantized values with relaxed tolerance. + if group_size is None: + a_deq = a * ref_scales.view(-1, 1) + b_deq = b * ops_scales.view(-1, 1) + else: + a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1) + b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1) + # NOTE: It is possible that some future test cases trigger this + # max diff due to precision issues. If such an error is + # encountered, it's recommended to inspect the differences between + # all corresponding elements from each tensor (e.g. by looping over + # them) and checking how many the max diff error shows up on (just + # a few bad elements should still be considered acceptable). + ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2) assert ok if add_residual: assert torch.allclose(ref_residual, ops_residual) diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index fde09710b5d..6e546f154c2 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/tests/kernels/core/test_layernorm.py @@ -5,7 +5,7 @@ import pytest import torch from tests.kernels.quant_utils import FP8_DTYPE -from tests.kernels.utils import opcheck +from tests.kernels.utils import fp8_ulp_distance, opcheck from vllm import ir from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.platforms import current_platform @@ -204,12 +204,22 @@ def test_fused_rms_norm_quant( (out_quant_fused, x, weight, quant_scale_t, 1e-6), ) - torch.testing.assert_close( - out_quant.to(dtype=torch.float32), - out_quant_fused.to(dtype=torch.float32), - atol=1e-3, - rtol=1e-3, - ) + if current_platform.is_rocm(): + # Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie; + # tolerate a tiny number of isolated fp8 outliers on ROCm. + ulp = fp8_ulp_distance(out_quant, out_quant_fused) + max_outliers = ulp.numel() // 100_000 + 8 + num_outliers = int((ulp > 0).sum().item()) + assert num_outliers <= max_outliers, ( + f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})" + ) + else: + torch.testing.assert_close( + out_quant.to(dtype=torch.float32), + out_quant_fused.to(dtype=torch.float32), + atol=1e-3, + rtol=1e-3, + ) @torch.inference_mode() diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index d2919185519..ed163a0472a 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -19,6 +19,7 @@ The kernel is imported via import pytest import torch +from tests.kernels.utils import bf16_ulp_distance, fp8_ulp_distance from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -160,35 +161,6 @@ def _call_fused( ) -def _bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """Representable-step distance between two bf16 tensors. - - Reinterprets the bf16 bit patterns under the IEEE-754 total ordering so - that adjacent representable values differ by exactly 1. - """ - - def key(t: torch.Tensor) -> torch.Tensor: - u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF - return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000) - - return (key(a) - key(b)).abs() - - -def _fp8_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """Representable-step distance between two 8-bit fp8 tensors. - - Reinterprets the fp8 bytes under a sign-magnitude total ordering so that - adjacent representable values differ by exactly 1. Inputs must already share - the same fp8 encoding (e.g. both FP8_STORE_DTYPE). - """ - - def key(t: torch.Tensor) -> torch.Tensor: - u = t.contiguous().view(torch.uint8).to(torch.int64) - return torch.where(u >= 0x80, 0xFF - u, u + 0x80) - - return (key(a) - key(b)).abs() - - def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor: """Reinterpret a float8_e4m3fn-typed kernel output under the real (FNUZ on gfx942) encoding the kernel actually wrote, without touching the bytes.""" @@ -235,7 +207,7 @@ def _assert_kv_cache_parity( rec_fused[:, :NOPE_DIM], rec_ref[:, :NOPE_DIM], rtol=0, atol=0 ) max_ulp = int( - _bf16_ulp_distance(rec_fused[:, NOPE_DIM:], rec_ref[:, NOPE_DIM:]).max().item() + bf16_ulp_distance(rec_fused[:, NOPE_DIM:], rec_ref[:, NOPE_DIM:]).max().item() ) assert max_ulp <= 1, f"RoPE bf16 region differs by {max_ulp} ULP (>1)" @@ -709,7 +681,7 @@ def test_full_cache_per_tensor_fp8_matches_reference( # reduction and RoPE rotation can land the kernel and the torch reference on # opposite sides of an fp8 round-to-nearest tie, so allow <=1 fp8 ULP. q_fused = _as_stored_fp8(q_fp8_fused) - q_max_ulp = int(_fp8_ulp_distance(q_fused, q_fp8_ref).max().item()) + q_max_ulp = int(fp8_ulp_distance(q_fused, q_fp8_ref).max().item()) assert q_max_ulp <= 1, f"Q fp8 differs by {q_max_ulp} ULP (>1)" # K-cache NoPE region [0, NOPE_DIM) is a deterministic per-tensor fp8 quant @@ -723,7 +695,7 @@ def test_full_cache_per_tensor_fp8_matches_reference( atol=0, ) k_max_ulp = int( - _fp8_ulp_distance(k_fused[..., NOPE_DIM:], k_cache_ref[..., NOPE_DIM:]) + fp8_ulp_distance(k_fused[..., NOPE_DIM:], k_cache_ref[..., NOPE_DIM:]) .max() .item() ) diff --git a/tests/kernels/utils.py b/tests/kernels/utils.py index 12ff3830c21..cc1d1bbf88d 100644 --- a/tests/kernels/utils.py +++ b/tests/kernels/utils.py @@ -809,6 +809,35 @@ def fp8_allclose( ) +def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Representable-step distance between two bf16 tensors. + + Reinterprets the bf16 bit patterns under the IEEE-754 total ordering so + that adjacent representable values differ by exactly 1. + """ + + def key(t: torch.Tensor) -> torch.Tensor: + u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF + return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000) + + return (key(a) - key(b)).abs() + + +def fp8_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Representable-step distance between two 8-bit fp8 tensors. + + Reinterprets the fp8 bytes under a sign-magnitude total ordering so that + adjacent representable values differ by exactly 1. Inputs must already share + the same fp8 encoding (e.g. both FP8_STORE_DTYPE). + """ + + def key(t: torch.Tensor) -> torch.Tensor: + u = t.contiguous().view(torch.uint8).to(torch.int64) + return torch.where(u >= 0x80, 0xFF - u, u + 0x80) + + return (key(a) - key(b)).abs() + + # Marlin MoE test utils From 00e045b7c7b82599f626779e111233abd4d0a64e Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Fri, 26 Jun 2026 22:45:23 -0400 Subject: [PATCH 064/138] [ROCm][CI TG] refactor and fix deepep_moe test group (#46758) Signed-off-by: Divakar Verma --- tests/kernels/moe/test_deepep_moe.py | 95 ++++++++++++++++++---------- tests/kernels/moe/test_ocp_mx_moe.py | 24 +------ tests/kernels/moe/utils.py | 23 +++++++ 3 files changed, 84 insertions(+), 58 deletions(-) diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 4080ca18459..8d12e2888d0 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -11,7 +11,7 @@ import torch.distributed from torch.distributed import ProcessGroup import vllm.envs as envs -from tests.kernels.moe.utils import make_dummy_moe_config +from tests.kernels.moe.utils import check_accuracy, make_dummy_moe_config from vllm import _custom_ops as ops from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.activation import SiluAndMul @@ -227,39 +227,44 @@ def deep_ep_moe_impl( out_hidden_states = torch.empty_like(test_tensors.rank_tokens) total_num_tokens = test_tensors.rank_tokens.size(0) + quant_config = FusedMoEQuantConfig.make( + q_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + per_act_token_quant=per_act_token_quant, + a1_scale=test_tensors.rank_token_scales, + ) + + # Build the kernel (and its DeepEP buffer) once and reuse it across chunks. + # Re-creating it per chunk re-inits rocSHMEM, which only allows one + # allocation per process on ROCm. The buffer is sized by max_tokens_per_rank + # so it is valid for every chunk (mirrors production's cached all2all handle). + mk: FusedMoEKernel = make_modular_kernel( + pg, + pgi, + low_latency_mode, + hidden_size, + dp_size, + num_experts, + num_local_experts, + q_dtype, + use_fp8_dispatch, + quant_config, + ) + def process_chunk(chunk_start, chunk_end, skip_result_store=False): rank_tokens_chunk = test_tensors.rank_tokens[chunk_start:chunk_end] topk_weights_chunk = test_tensors.topk_weights[chunk_start:chunk_end] topk_chunk = test_tensors.topk[chunk_start:chunk_end] - rank_token_scales_chunk = test_tensors.rank_token_scales - if ( - rank_token_scales_chunk is not None - and rank_token_scales_chunk.size(0) == total_num_tokens - ): - # per act token - rank_token_scales_chunk = rank_token_scales_chunk[chunk_start:chunk_end] - quant_config = FusedMoEQuantConfig.make( - q_dtype, - w1_scale=w1_scale, - w2_scale=w2_scale, - per_act_token_quant=per_act_token_quant, - a1_scale=rank_token_scales_chunk, - ) - - # Make modular kernel - mk: FusedMoEKernel = make_modular_kernel( - pg, - pgi, - low_latency_mode, - hidden_size, - dp_size, - num_experts, - num_local_experts, - q_dtype, - use_fp8_dispatch, - quant_config, - ) + if low_latency_mode: + # Reusing one buffer leaves it dirty; the low-latency kernels need + # the zero-initialized regions reset before each dispatch. + mk.prepare_finalize.buffer.clean_low_latency_buffer( + MAX_TOKENS_PER_RANK, + hidden_size, + num_experts, + ) out = mk.apply( hidden_states=rank_tokens_chunk, @@ -350,6 +355,28 @@ def torch_moe_impl( return out +def assert_deepep_close( + expected: torch.Tensor, + actual: torch.Tensor, + k: int, + use_fp8_dispatch: bool, +) -> None: + if use_fp8_dispatch and current_platform.is_fp8_fnuz(): + # ROCm e4m3fnuz rounds differently than the reference quant, + # so DeepEP's fp8 dispatch can yield a few outliers even with + # a correct kernel; allow a small fraction of mismatches here. + atol = rtol = 1.5e-1 + check_accuracy(expected, actual, atol=atol, rtol=rtol, percent=0.95) + return + + torch.testing.assert_close( + expected, + actual, + atol=6e-2, + rtol=6e-2, + ) + + def _deep_ep_moe( pgi: ProcessGroupInfo, low_latency_mode: bool, @@ -362,6 +389,9 @@ def _deep_ep_moe( use_fp8_dispatch: bool, per_act_token_quant: bool, ): + # Set seed in worker process for deterministic tensor generation. + set_random_seed(7) + device = torch.device(f"cuda:{pgi.local_rank}") init_workspace_manager(device) @@ -426,12 +456,7 @@ def _deep_ep_moe( per_act_token_quant, ) - torch.testing.assert_close( - torch_combined, - deepep_combined, - atol=6e-2, - rtol=6e-2, - ) + assert_deepep_close(torch_combined, deepep_combined, config.k, use_fp8_dispatch) MNKs = [ diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index e768947b269..a96e47fe439 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -9,6 +9,7 @@ import pytest import torch from packaging import version +from tests.kernels.moe.utils import check_accuracy from vllm._aiter_ops import is_aiter_found from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer @@ -515,29 +516,6 @@ def tg_mxfp4_moe( return tg_result -def check_accuracy(a, b, atol, rtol, percent): - """Allow a mismatch percentage of 1 - percent.""" - if torch.any(torch.isnan(a)): - raise Exception("NaN in reference output") - if torch.any(torch.isnan(b)): - raise Exception("NaN in actual output") - if torch.any(torch.isinf(a)): - raise Exception("Inf in reference output") - if torch.any(torch.isinf(b)): - raise Exception("Inf in actual output") - assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}" - - left = torch.abs(a - b) - right = atol + rtol * torch.abs(b) - count = torch.sum(left > right) - mismatch_percent = count / a.numel() - if mismatch_percent > 1 - percent: - raise Exception( - f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " - f"(threshold: {1 - percent:.4f})" - ) - - @pytest.mark.parametrize("topk", [1, 4]) @pytest.mark.parametrize("num_experts", [32, 128]) @pytest.mark.parametrize("num_tokens", [1, 128, 1024]) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4899de44a81..3f3bcebd11e 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -653,3 +653,26 @@ def make_shared_experts( return make_shared_experts_with_weights( N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype ) + + +def check_accuracy(a, b, atol, rtol, percent): + """Allow a mismatch percentage of 1 - percent.""" + if torch.any(torch.isnan(a)): + raise Exception("NaN in reference output") + if torch.any(torch.isnan(b)): + raise Exception("NaN in actual output") + if torch.any(torch.isinf(a)): + raise Exception("Inf in reference output") + if torch.any(torch.isinf(b)): + raise Exception("Inf in actual output") + assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}" + + left = torch.abs(a - b) + right = atol + rtol * torch.abs(b) + count = torch.sum(left > right) + mismatch_percent = count / a.numel() + if mismatch_percent > 1 - percent: + raise Exception( + f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " + f"(threshold: {1 - percent:.4f})" + ) From ddd3855a28a561a5bb54d380c6e6b8b1e883cc4a Mon Sep 17 00:00:00 2001 From: Cheng Jiang Date: Sat, 27 Jun 2026 11:18:07 +0800 Subject: [PATCH 065/138] [MoE Backend] add HPC-Ops MoE backend (#45924) Signed-off-by: chengvjiang Co-authored-by: chengvjiang Co-authored-by: youkaichao --- docs/design/moe_kernel_features.md | 1 + vllm/config/kernel.py | 2 + .../layers/fused_moe/hpc_moe.py | 211 +++++++++++++++++ .../layers/fused_moe/oracle/fp8.py | 18 +- vllm/utils/hpc.py | 223 ++++++++++++++++++ 5 files changed, 453 insertions(+), 2 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/hpc_moe.py create mode 100644 vllm/utils/hpc.py diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 279ab2d0d6f..d49790e833a 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -89,6 +89,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],
[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] | | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | +| hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.experts.hpc.HPCExperts] | | rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index e9f41c538ad..770daad1cef 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -134,6 +134,7 @@ MoEBackend = Literal[ "triton_unfused", "aiter", "flydsl", + "hpc", "emulation", ] @@ -189,6 +190,7 @@ class KernelConfig: - "triton_unfused": Use Triton unfused MoE kernels - "aiter": Use AMD AITer kernels (ROCm only) - "flydsl": Use AMD FlyDSL kernels (ROCm only) + - "hpc": Use HPC kernels (FP8 and Hopper only) - "emulation": use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations. """ diff --git a/vllm/model_executor/layers/fused_moe/hpc_moe.py b/vllm/model_executor/layers/fused_moe/hpc_moe.py new file mode 100644 index 00000000000..6c7e20fa3c1 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/hpc_moe.py @@ -0,0 +1,211 @@ +# 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.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform +from vllm.utils.hpc import has_hpc, hpc_fuse_moe, hpc_fuse_moe_blockwise + +logger = init_logger(__name__) + + +class HPCExperts(mk.FusedMoEExpertsModular): + """MoE implementation powered by [HPC](https://github.com/Tencent/hpc-ops). + + Only supported on NVIDIA Hopper GPUs (e.g. H20, H200), and currently limited to + FP8 models such as Hy3-FP8, Qwen3-235B-A22B-FP8, etc. + """ + + def __init__( + self, + moe_config: mk.FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + + assert quant_config.weight_quant_dtype in (torch.float8_e4m3fn,), ( + "Only fp8 quantization is currently supported." + ) + + self.device = moe_config.device + self.num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + self.ep_size = moe_config.moe_parallel_config.ep_size + self.tp_rank = moe_config.moe_parallel_config.tp_rank + self.tp_size = moe_config.moe_parallel_config.tp_size + self.out_dtype = moe_config.in_dtype + + @property + def expects_unquantized_inputs(self) -> bool: + return False + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return ( + p.is_cuda() + and (p.is_device_capability(90) or p.is_device_capability_family(100)) + and has_hpc() + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + scheme = (weight_key, activation_key) + # The following are supported by HPCExperts: + return scheme in [ + # fp8 static per-tensor on 9.0+ + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + ] + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in [ + MoEActivation.SILU, + ] + + @staticmethod + def _supports_shape(hidden_dim: int) -> bool: + # HPC fused MoE kernels process hidden_size in blocks of 128: + # block-wise fp8 requires hidden_size % 128 == 0 (per-128 quant), and + # the group GEMM tiles N by 128. Require 128-alignment to cover all + # code paths. + return hidden_dim % 128 == 0 + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def supports_chunking(self) -> bool: + # This refers to TP chunking; DP chunking is handled separately. + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + 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, ...]]: + # We use global_num_experts due to how moe_align_block_size handles + # expert_maps. + """ + Compute the shapes for the temporary and final outputs of the two gemms + and activation in the fused expert function. Since the gemms are + independent, the workspace for the first gemm can be shared with the + workspace for the last gemm. + + Returns a tuple of: + - workspace13 shape tuple: must be large enough to hold the + result of either expert gemm. + - workspace2 shape tuple: must be large enough to hold the + result of the activation function. + - output shape tuple: must be exact size of the final gemm output. + - Workspace type: The dtype to use for the workspace tensors. + - Note: in order for activation chunking to work, the first dimension + of each tuple must be the number of tokens. + """ + workspace1 = (M, K) + workspace2 = (0,) + output_shape = (M, K) + # The workspace is determined by `aq`, since it comes after any + # potential communication op and is involved in the expert computation. + return (workspace1, workspace2, output_shape) + + 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 | None, + ): + assert self._supports_activation(activation), f"{activation=} not supported" + assert self.quant_config.w1_scale is not None, ( + "w13_weight_scale must be provided" + ) + assert self.quant_config.w2_scale is not None, ( + "w2_weight_scale must be provided" + ) + + if self.quant_config.is_block_quantized: + hpc_fuse_moe_blockwise( + x=hidden_states, + x_scale=a1q_scale, + gate_up_weight=w1, + gate_up_weight_scale=self.quant_config.w1_scale, + down_weight=w2, + down_weight_scale=self.quant_config.w2_scale, + topk_ids=topk_ids, + topk_scale=topk_weights, + rank_ep=self.ep_rank, + num_expert_total=global_num_experts, + output=output, + ) + else: + assert self.quant_config.a1_scale is not None, ( + "w13_input_scale must be provided" + ) + assert self.quant_config.a2_scale is not None, ( + "w2_input_scale must be provided" + ) + hpc_fuse_moe( + x=hidden_states, + gate_up_weight=w1, + down_weight=w2, + gate_up_scale=self.quant_config.g1_alphas, + down_scale=self.quant_config.g2_alphas, + act_and_mul_scale=self.quant_config.a2_gscale, + topk_ids=topk_ids, + topk_scale=topk_weights, + rank_ep=self.ep_rank, + num_expert_total=global_num_experts, + output=output, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 1b5030b1909..9f930f1d58c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -52,6 +52,7 @@ class Fp8MoeBackend(Enum): BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS" XPU = "XPU" CPU = "CPU" + HPC = "HPC" # Dequantize-to-BF16 emulation for MXFP8 on devices without a native # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here. EMULATION = "EMULATION" @@ -85,6 +86,7 @@ def _get_priority_backends( Fp8MoeBackend.BATCHED_TRITON, Fp8MoeBackend.XPU, Fp8MoeBackend.CPU, + Fp8MoeBackend.HPC, ] def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None: @@ -216,6 +218,13 @@ def backend_to_kernel_cls( return [CPUExpertsFp8] + elif backend == Fp8MoeBackend.HPC: + from vllm.model_executor.layers.fused_moe.hpc_moe import ( + HPCExperts, + ) + + return [HPCExperts] + else: raise ValueError(f"Unknown FP8 MoE backend: {backend.value}") @@ -230,6 +239,7 @@ def map_fp8_backend(runner_backend: MoEBackend) -> Fp8MoeBackend: "flashinfer_cutlass": Fp8MoeBackend.FLASHINFER_CUTLASS, "marlin": Fp8MoeBackend.MARLIN, "aiter": Fp8MoeBackend.AITER, + "hpc": Fp8MoeBackend.HPC, } if backend := mapping.get(runner_backend): return backend @@ -470,6 +480,7 @@ def convert_to_fp8_moe_kernel_format( Fp8MoeBackend.VLLM_CUTLASS, Fp8MoeBackend.BATCHED_VLLM_CUTLASS, Fp8MoeBackend.XPU, + Fp8MoeBackend.HPC, # EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes # the MXFP8 weights as-is โ€” neither needs a load-time layout change. Fp8MoeBackend.EMULATION, @@ -521,9 +532,12 @@ def make_fp8_moe_quant_config( gemm1_clamp_limit=swiglu_limit, ) - # Flashinfer CUTLASS per-tensor uses single dq scale + # Flashinfer CUTLASS or HPC per-tensor uses single dq scale # (alpha = w_scale * a_scale) and inverse a2 scale. - if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS and block_shape is None: + if ( + fp8_backend in [Fp8MoeBackend.FLASHINFER_CUTLASS, Fp8MoeBackend.HPC] + and block_shape is None + ): assert a1_scale is not None and a2_scale is not None return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, diff --git a/vllm/utils/hpc.py b/vllm/utils/hpc.py new file mode 100644 index 00000000000..abe546c9165 --- /dev/null +++ b/vllm/utils/hpc.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility wrapper for HPC API changes. + +Users of vLLM should always import **only** these wrappers. +""" + +import functools +import importlib +import importlib.util + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@functools.cache +def has_hpc() -> bool: + """Return `True` if hpc package is available.""" + # Use find_spec to check if the module exists without importing it + # This avoids potential CUDA initialization side effects + if importlib.util.find_spec("hpc") is None: + logger.warning_once( + "HPC attention requires the hpc module to be installed. " + "Please install it from https://github.com/Tencent/hpc-ops" + ) + return False + return True + + +# Remove 'torch._library.custom_ops': +# The output of this custom operator (1) must not also be an input to +# this custom operator and (2) may not alias any inputs to this custom +# operator or other returns. The most common way to trigger this error +# is if we have y = custom_op(x) and y and x are the same Tensor. +# Please instead return a clone of the offending output tensor(s) (e.g. +# return x.clone()) or refactor the custom operator to not return y. +# @torch.library.custom_op( +# "vllm::fuse_moe_impl", +# mutates_args=[], +# device_types="cuda", +# ) +def fuse_moe_impl( + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_scale: torch.Tensor, + down_scale: torch.Tensor, + act_and_mul_scale: torch.Tensor, + topk_ids: torch.Tensor, + topk_scale: torch.Tensor, + rank_ep: int, + num_expert_total: int, + use_bf16_mul: bool = True, + shared_output: torch.Tensor = None, + output: torch.Tensor = None, +) -> torch.Tensor: + from hpc import fuse_moe as fuse_moe_ + + return fuse_moe_( + x, + gate_up_weight, + down_weight, + gate_up_scale, + down_scale, + act_and_mul_scale, + topk_ids, + topk_scale, + rank_ep, + num_expert_total, + use_bf16_mul, + shared_output, + output=output, + ) + + +# @torch.library.register_fake( +# "vllm::fuse_moe_impl", +# ) +def fuse_moe_impl_fake( + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_scale: torch.Tensor, + down_scale: torch.Tensor, + act_and_mul_scale: torch.Tensor, + topk_ids: torch.Tensor, + topk_scale: torch.Tensor, + rank_ep: int, + num_expert_total: int, + use_bf16_mul: bool = True, + shared_output: torch.Tensor = None, + output: torch.Tensor = None, +) -> torch.Tensor: + return torch.empty_like(x) + + +def hpc_fuse_moe( + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_scale: torch.Tensor, + down_scale: torch.Tensor, + act_and_mul_scale: torch.Tensor, + topk_ids: torch.Tensor, + topk_scale: torch.Tensor, + rank_ep: int, + num_expert_total: int, + use_bf16_mul: bool = True, + shared_output: torch.Tensor = None, + output: torch.Tensor = None, +) -> torch.Tensor: + return fuse_moe_impl( + x, + gate_up_weight, + down_weight, + gate_up_scale, + down_scale, + act_and_mul_scale, + topk_ids, + topk_scale, + rank_ep, + num_expert_total, + use_bf16_mul, + shared_output, + output=output, + ) + + +# @torch.library.custom_op( +# "vllm::fuse_moe_blockwise_impl", +# mutates_args=[], +# device_types="cuda", +# ) +def fuse_moe_blockwise_impl( + x: torch.Tensor, + x_scale: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + down_weight: torch.Tensor, + down_weight_scale: torch.Tensor, + topk_ids: torch.Tensor, + topk_scale: torch.Tensor, + rank_ep: int, + num_expert_total: int, + shared_output: torch.Tensor = None, + output: torch.Tensor = None, +) -> torch.Tensor: + from hpc import fuse_moe_blockwise as fuse_moe_blockwise_ + + return fuse_moe_blockwise_( + x, + x_scale, + gate_up_weight, + gate_up_weight_scale, + down_weight, + down_weight_scale, + topk_ids, + topk_scale, + rank_ep, + num_expert_total, + shared_output, + output=output, + ) + + +# @torch.library.register_fake( +# "vllm::fuse_moe_blockwise_impl", +# ) +def fuse_moe_blockwise_impl_fake( + x: torch.Tensor, + x_scale: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + down_weight: torch.Tensor, + down_weight_scale: torch.Tensor, + topk_ids: torch.Tensor, + topk_scale: torch.Tensor, + rank_ep: int, + num_expert_total: int, + shared_output: torch.Tensor = None, + output: torch.Tensor = None, +) -> torch.Tensor: + return torch.empty_like(x) + + +def hpc_fuse_moe_blockwise( + x: torch.Tensor, + x_scale: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + down_weight: torch.Tensor, + down_weight_scale: torch.Tensor, + topk_ids: torch.Tensor, + topk_scale: torch.Tensor, + rank_ep: int, + num_expert_total: int, + shared_output: torch.Tensor = None, + output: torch.Tensor = None, +) -> torch.Tensor: + return fuse_moe_blockwise_impl( + x, + x_scale, + gate_up_weight, + gate_up_weight_scale, + down_weight, + down_weight_scale, + topk_ids, + topk_scale, + rank_ep, + num_expert_total, + shared_output, + output=output, + ) + + +__all__ = [ + "has_hpc", + "hpc_fuse_moe", + "hpc_fuse_moe_blockwise", +] From 68ee8300a047db78fb52bac477daaaac7be11216 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Sat, 27 Jun 2026 00:38:13 -0400 Subject: [PATCH 066/138] [ROCm][CI]Fix test_concat_and_cache_mla_rope_fused on ROCm (#46409) Signed-off-by: Divakar Verma --- .../test_rotary_embedding_mla_cache_fused.py | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py index 181f10f314e..289267b6a4c 100644 --- a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py +++ b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py @@ -17,6 +17,36 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +@pytest.fixture +def default_vllm_config(monkeypatch): + """Enable the AITER triton rope on ROCm for fp16-consistent numerics. + + The fused CUDA kernel runs native fp16 while forward_native upcasts to + fp32, so on ROCm we route through the AITER triton rope (+rotary_embedding) + to match. Its env gates are cached at import, hence refresh_env_variables(). + """ + from vllm._aiter_ops import rocm_aiter_ops + from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config + + is_rocm = current_platform.is_rocm() + if is_rocm: + config = VllmConfig( + compilation_config=CompilationConfig(custom_ops=["+rotary_embedding"]) + ) + else: + config = VllmConfig() + try: + with monkeypatch.context() as m, set_current_vllm_config(config): + if is_rocm: + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "1") + rocm_aiter_ops.refresh_env_variables() + yield config + finally: + if is_rocm: + rocm_aiter_ops.refresh_env_variables() + + @pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float]) @pytest.mark.parametrize("is_neox_style", [False, True]) @pytest.mark.parametrize("seq_len", [11, 42]) @@ -151,6 +181,10 @@ def test_concat_and_cache_mla_rope_fused( kv_cache_scale, ) + # ROCm neox-style Triton FMA diverges slightly from the fused kernel, so + # relax the affected tolerance: rtol for fp8 (one e4m3 ULP ~12.5%) and atol + # otherwise (bounded ~6e-4). Other paths use the CUDA defaults. + rocm_neox = current_platform.is_rocm() and is_neox_style if kv_cache_dtype == "fp8": result_temp = torch.empty_like(kv_cache, dtype=torch.float16) ops.convert_fp8( @@ -163,7 +197,11 @@ def test_concat_and_cache_mla_rope_fused( ops.convert_fp8( expected_temp, ref_kv_cache, kv_cache_scale.item(), kv_dtype=kv_cache_dtype ) - torch.testing.assert_close(result_temp, expected_temp, atol=0.001, rtol=0.1) + torch.testing.assert_close( + result_temp, expected_temp, atol=0.001, rtol=0.15 if rocm_neox else 0.1 + ) + elif rocm_neox: + torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3) else: torch.testing.assert_close(kv_cache, ref_kv_cache) From d706dec904e89e4067efd9535b96705aa61f3935 Mon Sep 17 00:00:00 2001 From: JasonCohere Date: Sat, 27 Jun 2026 06:15:06 +0100 Subject: [PATCH 067/138] fix: Correct reasoning-end detection for prompt history (#44551) Signed-off-by: jwzheng96 Signed-off-by: JianweiZheng <32029023+jwzheng96@users.noreply.github.com> Signed-off-by: Jason Ozuzu Signed-off-by: walterbm Co-authored-by: JianweiZheng <32029023+jwzheng96@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: walterbm Co-authored-by: Walter Beller-Morales Co-authored-by: Flora Feng <4florafeng@gmail.com> --- requirements/test/cuda.in | 1 + requirements/test/cuda.txt | 2 + .../test_cohere_command_reasoning_parser.py | 625 ++++++++++++++++++ .../cohere_command_reasoning_parser.py | 18 +- 4 files changed, 645 insertions(+), 1 deletion(-) create mode 100644 tests/reasoning/test_cohere_command_reasoning_parser.py diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 03218c75e1f..12a40716392 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -73,6 +73,7 @@ gpt-oss >= 0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test +cohere_melody>=0.9.0 # required for cohere command reasoning parser test # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. # Older versions are in conflict with teerratorch requirements. diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 1a9fe6f16a0..f504c69c48f 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -101,6 +101,8 @@ click==8.1.7 # schemathesis # typer # uvicorn +cohere-melody==0.9.0 + # via -r requirements/test/cuda.in colorama==0.4.6 # via # perceptron diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py new file mode 100644 index 00000000000..a6524ca7072 --- /dev/null +++ b/tests/reasoning/test_cohere_command_reasoning_parser.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import json +from collections import UserDict +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + JsonSchemaResponseFormat, + ResponseFormat, + StructuralTagResponseFormat, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.reasoning.cohere_command_reasoning_parser import ( + CohereCommand3ReasoningParser, + CohereCommand4ReasoningParser, + _has_effective_tools, + _response_format_type, + _schema_dict_from_structured_outputs, + convert_schema_to_structural_tags, +) +from vllm.sampling_params import StructuredOutputsParams + + +@dataclass +class ExpectedToolCall: + id: str + name: str + arguments: dict + + +@dataclass +class ReasoningCase: + parser_cls: Any + model_output: str + expected_reasoning: str | None + expected_content: str | None + expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) + + +REASONING_CASES = [ + pytest.param( + ReasoningCase( + parser_cls=CohereCommand3ReasoningParser, + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_reasoning="i will call foo with query1", + expected_content="""\ +<|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_tool_calls=[ + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), + ], + ), + id="cmd3-single_tool_call", + ), + pytest.param( + ReasoningCase( + parser_cls=CohereCommand4ReasoningParser, + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_reasoning="i will call foo with query1", + expected_content="""\ +<|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_tool_calls=[ + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), + ], + ), + id="cmd4-single_tool_call", + ), + pytest.param( + ReasoningCase( + parser_cls=CohereCommand3ReasoningParser, + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: ๐ŸŒˆ<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: ๐ŸŒˆ", + expected_content="foo bar", + ), + id="cmd3-citations_with_emoji", + ), + pytest.param( + ReasoningCase( + parser_cls=CohereCommand4ReasoningParser, + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: ๐ŸŒˆ<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: ๐ŸŒˆ", + expected_content="foo bar", + ), + id="cmd4-citations_with_emoji", + ), +] + + +class MockCohereTokenizer: + """Minimal byte-level stand-in for the Cohere tokenizer. + + ``encode``/``decode`` round-trip through UTF-8 bytes so splitting a + multi-byte character (e.g. an emoji) across "tokens" reproduces the + trailing U+FFFD buffering that real streaming exhibits. Cohere special + tokens map to distinct synthetic ids; everything else shares a default id. + ``adjust_request`` only needs the token ids, not real tokenization. + """ + + _SPECIAL_TOKEN_IDS = { + "<|START_THINKING|>": -1, + "<|END_THINKING|>": -2, + "<|CHATBOT_TOKEN|>": -3, + } + + def convert_tokens_to_ids(self, token: str) -> int: + return self._SPECIAL_TOKEN_IDS.get(token, 0) + + def get_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return list(text.encode("utf-8")) + + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: + return bytes(ids).decode("utf-8", errors="replace") + + +@pytest.fixture(scope="module") +def tokenizer() -> MockCohereTokenizer: + return MockCohereTokenizer() + + +@pytest.fixture +def request_obj(): + return ChatCompletionRequest(messages=[], model="test-model") + + +REPLACEMENT_CHAR = "\ufffd" + + +def _token_deltas(tokenizer, text: str) -> list[str]: + """Progressively decode the token sequence and return per-step string + deltas. Incomplete multi-byte sequences (trailing U+FFFD) are buffered + until the next token completes them, matching real streaming behaviour.""" + ids = tokenizer.encode(text, add_special_tokens=False) + deltas: list[str] = [] + prev = "" + for i in range(1, len(ids) + 1): + current = tokenizer.decode(ids[:i], skip_special_tokens=False) + if current.endswith(REPLACEMENT_CHAR): + continue + delta = current[len(prev) :] + if delta: + deltas.append(delta) + prev = current + return deltas + + +@pytest.mark.parametrize("case", REASONING_CASES) +class TestExtractReasoning: + def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase): + parser = case.parser_cls(tokenizer) + reasoning, content = parser.extract_reasoning(case.model_output, request_obj) + + assert reasoning == case.expected_reasoning + assert content == case.expected_content + + def test_streaming(self, tokenizer, case: ReasoningCase): + parser = case.parser_cls(tokenizer) + token_strings = _token_deltas(tokenizer, case.model_output) + + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + tool_call_deltas: list[dict] = [] + + previous_text = "" + previous_token_ids: list[int] = [] + + for token_str in token_strings: + current_text = previous_text + token_str + current_token_ids = previous_token_ids + [0] + + delta = parser.extract_reasoning_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=token_str, + previous_token_ids=previous_token_ids, + current_token_ids=current_token_ids, + delta_token_ids=[0], + ) + if delta is not None: + if delta.reasoning is not None: + reasoning_parts.append(delta.reasoning) + if delta.content is not None: + content_parts.append(delta.content) + for tc in delta.tool_calls: + tool_call_deltas.append( + { + "id": tc.id, + "index": tc.index, + "name": tc.function.name if tc.function else None, + "arguments": ( + tc.function.arguments if tc.function else None + ), + } + ) + + previous_text = current_text + previous_token_ids = current_token_ids + + reasoning = "".join(reasoning_parts) if reasoning_parts else None + assert reasoning == case.expected_reasoning + + content = "".join(content_parts) if content_parts else None + if case.expected_tool_calls: + assert content is None or content == "" + else: + assert content == case.expected_content + + accumulated: dict[int, dict] = {} + for d in tool_call_deltas: + idx = d["index"] + if idx not in accumulated: + accumulated[idx] = {"id": "", "name": "", "arguments": ""} + if d["id"]: + accumulated[idx]["id"] = d["id"] + if d["name"]: + accumulated[idx]["name"] = d["name"] + if d["arguments"]: + accumulated[idx]["arguments"] += d["arguments"] + + assert len(accumulated) == len(case.expected_tool_calls) + for i, expected_tc in enumerate(case.expected_tool_calls): + tc = accumulated[i] + assert tc["id"] == expected_tc.id + assert tc["name"] == expected_tc.name + assert json.loads(tc["arguments"]) == expected_tc.arguments + + +class TestIsReasoningEnd: + @pytest.mark.parametrize( + "parser_cls", + [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser], + ids=["cmd3", "cmd4"], + ) + def test_is_reasoning_end(self, tokenizer, parser_cls): + parser = parser_cls(tokenizer) + start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") + end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") + content_ids = [99, 100] + + # Generation-only tokens have no chatbot marker, so the whole sequence + # is considered. + assert parser.is_reasoning_end([end_id]) + assert parser.is_reasoning_end([start_id, *content_ids, end_id]) + assert not parser.is_reasoning_end([start_id, *content_ids]) + + # Full prompt/history tokens are scoped to the latest chatbot marker, + # so stray thinking tokens from the preamble or previous turns are ignored. + assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids]) + assert parser.is_reasoning_end( + [start_id, end_id, chatbot_id, start_id, *content_ids, end_id] + ) + + +SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}} +SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}} +GET_WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, +} +VALID_STRUCTURAL_TAG = { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "tags": [ + { + "begin": "", + "content": {"type": "any_text"}, + "end": "", + } + ], + "triggers": [""], + }, +} + + +def _model_config(arch: str) -> SimpleNamespace: + return SimpleNamespace( + architecture=arch, + architectures=[arch], + hf_text_config=SimpleNamespace(architectures=[arch]), + ) + + +def _make_chat_request(**kwargs) -> ChatCompletionRequest: + data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"} + data.update(kwargs) + return ChatCompletionRequest.model_validate(data) + + +def _first_json_schema(tag_json: str) -> dict | None: + outer = json.loads(tag_json) + for t in (outer.get("format") or {}).get("tags") or []: + c = t.get("content") or {} + if c.get("type") == "json_schema": + js = c.get("json_schema") + return js if isinstance(js, dict) else None + return None + + +def _content_types(tag_json: str) -> set[str]: + outer = json.loads(tag_json) + out: set[str] = set() + for t in (outer.get("format") or {}).get("tags") or []: + ty = (t.get("content") or {}).get("type") + if isinstance(ty, str): + out.add(ty) + return out + + +@pytest.fixture(scope="module") +def parser(tokenizer: MockCohereTokenizer) -> CohereCommand4ReasoningParser: + """Parser configured with a supported Cohere architecture.""" + return CohereCommand4ReasoningParser( + tokenizer, + model_config=_model_config("Cohere2ForCausalLM"), + ) + + +@pytest.fixture(scope="module") +def parser_no_model_config( + tokenizer: MockCohereTokenizer, +) -> CohereCommand4ReasoningParser: + """Parser with no ``model_config`` (cannot resolve architecture).""" + return CohereCommand4ReasoningParser(tokenizer, model_config=None) + + +@pytest.fixture(scope="module") +def parser_unsupported_arch( + tokenizer: MockCohereTokenizer, +) -> CohereCommand4ReasoningParser: + """Parser configured with an architecture that has no structural tag style.""" + return CohereCommand4ReasoningParser( + tokenizer, + model_config=_model_config("LlamaForCausalLM"), + ) + + +class TestAdjustRequestPassthrough: + def test_structured_outputs_structural_tag_not_modified(self, parser) -> None: + tag = json.dumps(VALID_STRUCTURAL_TAG) + r = _make_chat_request(structured_outputs={"structural_tag": tag}) + o = parser.adjust_request(r) + assert o.structured_outputs.structural_tag == tag + + def test_response_format_structural_tag_short_circuit(self, parser) -> None: + # ``ChatCompletionRequest`` validates ``response_format`` as a union; + # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model). + rf = StructuralTagResponseFormat( + type="structural_tag", + format=VALID_STRUCTURAL_TAG["format"], + ) + r = _make_chat_request(response_format=rf) + o = parser.adjust_request(r) + assert _response_format_type(o.response_format) == "structural_tag" + assert o.structured_outputs is None + + +class TestAdjustRequestNoOp: + def test_no_schema_no_tools(self, parser) -> None: + o = parser.adjust_request(_make_chat_request()) + assert o.structured_outputs is None + assert o.response_format is None + + def test_no_model_config(self, parser_no_model_config) -> None: + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + ) + o = parser_no_model_config.adjust_request(r) + assert o.response_format is not None + assert o.structured_outputs is None + + +class TestAdjustRequestUnsupportedArchitecture: + def test_json_schema_raises(self, parser_unsupported_arch) -> None: + inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + ) + with pytest.raises(ValueError, match="does not support"): + parser_unsupported_arch.adjust_request(r) + + +class TestAdjustRequestFoldFromResponseFormat: + @pytest.mark.parametrize( + "response_format, expected_schema", + [ + pytest.param( + ResponseFormat( + type="json_schema", + json_schema=JsonSchemaResponseFormat( + name="n", json_schema=SCHEMA_A + ), + ), + SCHEMA_A, + id="json_schema_pydantic", + ), + pytest.param( + { + "type": "json_schema", + "json_schema": {"name": "n", "schema": SCHEMA_A}, + }, + SCHEMA_A, + id="json_schema_dict", + ), + pytest.param( + {"type": "json_object"}, + {"type": "object"}, + id="json_object", + ), + ], + ) + def test_response_format_cleared( + self, parser, response_format, expected_schema + ) -> None: + r = _make_chat_request(response_format=response_format) + o = parser.adjust_request(r) + assert o.response_format is None + assert ( + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema + ) + + +class TestHasEffectiveTools: + @pytest.mark.parametrize( + "tools, expected", + [ + pytest.param(None, False, id="none"), + pytest.param([], False, id="empty_list"), + pytest.param(" ", False, id="blank_str"), + pytest.param( + [{"type": "function", "function": {"name": "f"}}], + True, + id="non_empty_list", + ), + pytest.param('{"x": 1}', True, id="non_empty_str"), + ], + ) + def test_has_effective_tools(self, tools, expected) -> None: + assert _has_effective_tools(tools) is expected + + def test_convert_schema_json_only_with_empty_tools_list(self) -> None: + tag = convert_schema_to_structural_tags( + schema=SCHEMA_B, + tools=[], + model_architecture="Cohere2ForCausalLM", + ) + assert tag is not None + assert _first_json_schema(tag) == SCHEMA_B + + +class TestAdjustRequestFoldFromStructuredOutputs: + @pytest.mark.parametrize( + "structured_outputs, expected_schema", + [ + pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"), + pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"), + pytest.param( + {"json_object": True}, {"type": "object"}, id="json_object_flag" + ), + pytest.param( + StructuredOutputsParams(json=SCHEMA_B), + SCHEMA_B, + id="structured_outputs_dataclass", + ), + pytest.param( + {"json": {"name": "n", "schema": SCHEMA_A}}, + SCHEMA_A, + id="openai_wrapper_dict_unwrapped", + ), + ], + ) + def test_structured_outputs_folded( + self, parser, structured_outputs, expected_schema + ) -> None: + o = parser.adjust_request( + _make_chat_request(structured_outputs=structured_outputs), + ) + assert ( + _first_json_schema(o.structured_outputs.structural_tag) == expected_schema + ) + + def test_responses_request_default_empty_tools(self, parser) -> None: + """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``.""" + r = ResponsesRequest.model_validate( + { + "input": "hi", + "model": "m", + "structured_outputs": {"json": SCHEMA_B}, + } + ) + assert r.tools == [] + o = parser.adjust_request(r) + assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B + + def test_json_userdict_mapping_unwrapped(self) -> None: + inner = {"type": "object", "properties": {"u": {"type": "number"}}} + so = StructuredOutputsParams(json=UserDict(inner)) + assert _schema_dict_from_structured_outputs(so) == inner + + @pytest.mark.parametrize( + "json_value, match", + [ + pytest.param("{not json}", "valid JSON", id="invalid_json_string"), + pytest.param( + json.dumps(["a", "b"]), "JSON object", id="non_object_json_string" + ), + pytest.param(" ", "empty", id="empty_json_string"), + ], + ) + def test_structured_outputs_json_string_raises( + self, parser, json_value, match + ) -> None: + with pytest.raises(ValueError, match=match): + parser.adjust_request( + _make_chat_request(structured_outputs={"json": json_value}), + ) + + @pytest.mark.parametrize( + "construct", + [ + pytest.param( + lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}), + id="chat_completion_request", + ), + pytest.param( + lambda: StructuredOutputsParams(json=[1, 2, 3]), # type: ignore[arg-type] + id="structured_outputs_params", + ), + ], + ) + def test_json_wrong_type_raises(self, construct) -> None: + """Non-str / non-dict ``json`` fails at Pydantic validation.""" + with pytest.raises(ValidationError): + construct() + + +class TestAdjustRequestPrecedence: + def test_response_format_over_structured_outputs_json(self, parser) -> None: + s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}} + s_so = {"type": "object", "properties": {"so": {"type": "number"}}} + inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + structured_outputs={"json": s_so}, + ) + o = parser.adjust_request(r) + assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf + + +class TestAdjustRequestTextPlusStructuredOutputs: + def test_text_response_format_preserved(self, parser) -> None: + sch = {"type": "object", "properties": {"k": {"type": "string"}}} + r = _make_chat_request( + response_format=ResponseFormat(type="text"), + structured_outputs={"json": sch}, + ) + o = parser.adjust_request(r) + assert o.response_format is not None + assert o.response_format.type == "text" + assert _first_json_schema(o.structured_outputs.structural_tag) == sch + + +class TestAdjustRequestTools: + def test_tools_only_command_a_grammar(self, parser) -> None: + o = parser.adjust_request( + _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"), + ) + assert "grammar" in _content_types(o.structured_outputs.structural_tag) + + def test_tools_plus_json_schema_both_kinds(self, parser) -> None: + inner = JsonSchemaResponseFormat( + name="n", + json_schema={"type": "object", "properties": {"r": {"type": "string"}}}, + ) + r = _make_chat_request( + response_format=ResponseFormat(type="json_schema", json_schema=inner), + tools=[GET_WEATHER_TOOL], + tool_choice="auto", + ) + o = parser.adjust_request(r) + types = _content_types(o.structured_outputs.structural_tag) + assert "grammar" in types + assert "json_schema" in types diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 34066ef2d92..f0e7aed0b03 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -414,7 +414,9 @@ class BaseCohereCommandReasoningParser(ReasoningParser): **kwargs, ): super().__init__(tokenizer, *args, **kwargs) + self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") + self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") self.unary_opts = unary_opts self.melody_unary = PyFilter(unary_opts) self.melody_streaming = PyFilter(streaming_opts) @@ -478,7 +480,21 @@ class BaseCohereCommandReasoningParser(ReasoningParser): return content_ids def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - return any(tid == self.end_token_id for tid in reversed(input_ids)) + chatbot = self.chatbot_token_id + start = self.start_token_id + end = self.end_token_id + has_end_token = False + + for i in reversed(range(len(input_ids))): + tid = input_ids[i] + if tid == start: + return has_end_token + if tid == chatbot: + return False + if tid == end: + has_end_token = True + + return has_end_token def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest From 455f25aa13905e189b9268298c812f2048100f6f Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Sat, 27 Jun 2026 01:15:10 -0400 Subject: [PATCH 068/138] [CLI] Add flag to print TTFT and TPS in `vllm chat` (#46775) Signed-off-by: Benjamin Chislett --- docs/cli/README.md | 6 +++ vllm/entrypoints/cli/openai.py | 70 +++++++++++++++++++++++++++++----- 2 files changed, 67 insertions(+), 9 deletions(-) diff --git a/docs/cli/README.md b/docs/cli/README.md index 08e986a7463..43857704522 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -80,6 +80,9 @@ vllm chat --url http://{vllm-serve-host}:{vllm-serve-port}/v1 # Quick chat with a single prompt vllm chat --quick "hi" + +# Print TTFT and throughput statistics after each response +vllm chat --stats ``` See [vllm chat](./chat.md) for the full reference of all available arguments. @@ -97,6 +100,9 @@ vllm complete --url http://{vllm-serve-host}:{vllm-serve-port}/v1 # Quick complete with a single prompt vllm complete --quick "The future of AI is" + +# Print TTFT and throughput statistics after each response +vllm complete --stats ``` See [vllm complete](./complete.md) for the full reference of all available arguments. diff --git a/vllm/entrypoints/cli/openai.py b/vllm/entrypoints/cli/openai.py index 1c18b193d1c..c9869077c0a 100644 --- a/vllm/entrypoints/cli/openai.py +++ b/vllm/entrypoints/cli/openai.py @@ -5,6 +5,7 @@ import argparse import os import signal import sys +import time from typing import TYPE_CHECKING from openai import OpenAI @@ -44,25 +45,58 @@ def _interactive_cli(args: argparse.Namespace) -> tuple[str, OpenAI]: return model_name, openai_client -def _print_chat_stream(stream) -> str: +def _print_chat_stream(stream, stats: bool = False) -> str: output = "" + start = time.perf_counter() + ttft: float | None = None + completion_tokens = 0 for chunk in stream: + if chunk.usage is not None: + completion_tokens = chunk.usage.completion_tokens + if not chunk.choices: + continue delta = chunk.choices[0].delta if delta.content: + if ttft is None: + ttft = time.perf_counter() - start output += delta.content print(delta.content, end="", flush=True) print() + if stats: + _print_metrics(start, ttft, completion_tokens) return output -def _print_completion_stream(stream) -> str: +def _print_metrics(start: float, ttft: float | None, completion_tokens: int) -> None: + total_time = time.perf_counter() - start + if ttft is None or total_time <= 0: + return + print(f"{'TTFT:':<5} {ttft * 1000:.2f} ms") + print( + f"{'TPS:':<5} {completion_tokens / total_time:.2f} tokens/s " + f"({completion_tokens} tokens in {total_time:.2f}s)" + ) + + +def _print_completion_stream(stream, stats: bool = False) -> str: output = "" + start = time.perf_counter() + ttft: float | None = None + completion_tokens = 0 for chunk in stream: + if chunk.usage is not None: + completion_tokens = chunk.usage.completion_tokens + if not chunk.choices: + continue text = chunk.choices[0].text - if text is not None: + if text: + if ttft is None: + ttft = time.perf_counter() - start output += text print(text, end="", flush=True) print() + if stats: + _print_metrics(start, ttft, completion_tokens) return output @@ -127,18 +161,23 @@ class ChatCommand(CLISubcommand): def cmd(args: argparse.Namespace) -> None: model_name, client = _interactive_cli(args) system_prompt = args.system_prompt + stats = args.stats conversation: list[ChatCompletionMessageParam] = [] if system_prompt is not None: conversation.append({"role": "system", "content": system_prompt}) + create_kwargs = {"model": model_name, "stream": True} + if stats: + create_kwargs["stream_options"] = {"include_usage": True} + if args.quick: conversation.append({"role": "user", "content": args.quick}) stream = client.chat.completions.create( - model=model_name, messages=conversation, stream=True + messages=conversation, **create_kwargs ) - output = _print_chat_stream(stream) + output = _print_chat_stream(stream, stats) conversation.append({"role": "assistant", "content": output}) return @@ -151,9 +190,9 @@ class ChatCommand(CLISubcommand): conversation.append({"role": "user", "content": input_message}) stream = client.chat.completions.create( - model=model_name, messages=conversation, stream=True + messages=conversation, **create_kwargs ) - output = _print_chat_stream(stream) + output = _print_chat_stream(stream, stats) conversation.append({"role": "assistant", "content": output}) @staticmethod @@ -176,6 +215,11 @@ class ChatCommand(CLISubcommand): metavar="MESSAGE", help=("Send a single prompt as MESSAGE and print the response, then exit."), ) + parser.add_argument( + "--stats", + action="store_true", + help="Print TTFT and TPS statistics after each response.", + ) return parser def subparser_init( @@ -198,6 +242,7 @@ class CompleteCommand(CLISubcommand): @staticmethod def cmd(args: argparse.Namespace) -> None: model_name, client = _interactive_cli(args) + stats = args.stats kwargs = { "model": model_name, @@ -205,10 +250,12 @@ class CompleteCommand(CLISubcommand): } if args.max_tokens: kwargs["max_tokens"] = args.max_tokens + if stats: + kwargs["stream_options"] = {"include_usage": True} if args.quick: stream = client.completions.create(prompt=args.quick, **kwargs) - _print_completion_stream(stream) + _print_completion_stream(stream, stats) return print("Please enter prompt to complete:") @@ -218,7 +265,7 @@ class CompleteCommand(CLISubcommand): except EOFError: break stream = client.completions.create(prompt=input_prompt, **kwargs) - _print_completion_stream(stream) + _print_completion_stream(stream, stats) @staticmethod def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: @@ -236,6 +283,11 @@ class CompleteCommand(CLISubcommand): metavar="PROMPT", help="Send a single prompt and print the completion output, then exit.", ) + parser.add_argument( + "--stats", + action="store_true", + help="Print TTFT and TPS statistics after each response.", + ) return parser def subparser_init( From b588f66dc2982fe3228e0aee55b80387d31db2e4 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 27 Jun 2026 01:16:20 -0400 Subject: [PATCH 069/138] [GLM5.2 Perf] `fused_indexer_q_rope_quant` triton kernel, 1.9% ~ 3.3% E2E Throughput improvement. (#46862) Signed-off-by: yewentao256 --- .../layers/sparse_attn_indexer.py | 131 ++++++++++++++++++ vllm/model_executor/models/deepseek_v2.py | 36 +++++ 2 files changed, 167 insertions(+) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index fe2b268cde6..c1bc731ee62 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -11,7 +11,11 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, @@ -37,6 +41,133 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 MXFP4_BLOCK_SIZE = 32 +@triton.jit +def _fused_indexer_q_rope_quant_kernel( + positions, + q, + q_s0, + q_s1, + cos_sin_cache, + cos_sin_s0, + q_fp8, + q_fp8_s0, + q_fp8_s1, + weights, + weights_s0, + weights_s1, + weights_out, + weights_out_s0, + weights_out_s1, + softmax_scale, + head_scale, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + is_neox: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + offs32 = tl.arange(0, 32) + offs64 = tl.arange(0, 64) + + pos = tl.load(positions + token) + cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32) + sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32) + q_base = q + token * q_s0 + head * q_s1 + out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1 + + if is_neox: + # NeoX layout, x0 = q[0:32], x1 = q[32:64] + x0 = tl.load(q_base + offs32).to(tl.float32) + x1 = tl.load(q_base + 32 + offs32).to(tl.float32) + else: + # interleaved layout + # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...] + x0 = tl.load(q_base + offs32 * 2).to(tl.float32) + x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32) + r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32) + r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32) + amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1))) + + q_nope = tl.load(q_base + 64 + offs64).to(tl.float32) + amax = tl.maximum(amax, tl.max(tl.abs(q_nope))) + scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max) + # e8m0 format + q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) + + if is_neox: + tl.store( + out_base + offs32, + tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + 32 + offs32, + tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + else: + tl.store( + out_base + offs32 * 2, + tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + offs32 * 2 + 1, + tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + 64 + offs64, + tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + + weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32) + tl.store( + weights_out + token * weights_out_s0 + head * weights_out_s1, + weight * q_scale * softmax_scale * head_scale, + ) + + +def fused_indexer_q_rope_quant( + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float, + head_scale: float, + is_neox: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + assert current_platform.is_cuda() + assert q.dtype == torch.bfloat16 + assert q.shape[-1] == 128 + assert cos_sin_cache.shape[-1] == 64 + assert weights.shape == q.shape[:2] + + q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype()) + weights_out = torch.empty_like(weights, dtype=torch.float32) + fp8_min, fp8_max = get_fp8_min_max() + _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])]( + positions, + q, + q.stride(0), + q.stride(1), + cos_sin_cache, + cos_sin_cache.stride(0), + q_fp8, + q_fp8.stride(0), + q_fp8.stride(1), + weights, + weights.stride(0), + weights.stride(1), + weights_out, + weights_out.stride(0), + weights_out.stride(1), + softmax_scale, + head_scale, + fp8_min=fp8_min, + fp8_max=fp8_max, + is_neox=is_neox, + num_warps=1, + ) + return q_fp8, weights_out + + def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 814118f8a79..9b08ca9825e 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -73,6 +73,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.sparse_attn_indexer import ( SparseAttnIndexer, + fused_indexer_q_rope_quant, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -674,6 +675,13 @@ class Indexer(nn.Module): ) self.is_inplace_rope = is_inplace_rope + self.use_fused_indexer_q = ( + current_platform.is_cuda() + and self.quant_block_size == self.head_dim + and self.head_dim == 128 + and self.rope_dim == 64 + and self.scale_fmt is not None + ) def forward( self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb @@ -698,6 +706,34 @@ class Indexer(nn.Module): rotary_emb( positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1) ) + elif self.use_fused_indexer_q and q.dtype == torch.bfloat16: + # 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( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + + q_fp8, weights = fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + weights, + self.softmax_scale, + self.n_head**-0.5, + rotary_emb.is_neox_style, + ) + + # rotate only the MQA K + q_dummy = torch.empty_like(k_pe.unsqueeze(1)) + _, k_pe = rotary_emb(positions, q_dummy, k_pe.unsqueeze(1)) + k_pe = k_pe.reshape(-1, 1, self.rope_dim) + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + + return self.indexer_op(hidden_states, q_fp8, k, weights) else: q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 From 091d13976c1c246714bb2112dd2e208561dda6a3 Mon Sep 17 00:00:00 2001 From: peizhang56 Date: Fri, 26 Jun 2026 23:35:50 -0700 Subject: [PATCH 070/138] [ROCm][CI] Add TRITON_ATTN score absolute tolerance floor (#46891) Signed-off-by: pei.zhang Co-authored-by: Claude --- .../scoring/test_cross_encoder_online_vision.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py index e6b4d3f873e..c6663dbdff0 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py @@ -38,12 +38,17 @@ BACKEND_TOL: dict[str, float] = { "FLEX_ATTENTION": 0.045, # gfx950:~3.25%, gfx942:~1.10% } -# ROCm 7.2/gfx950 shows small absolute drift on the low text-vs-text -# probability even though larger scores remain well inside the relative -# tolerance. Keep the relative tolerances tight and add only a small floor. +# Some ROCm attention backends show small absolute drift on the low +# text-vs-text probability even though larger scores remain well inside the +# relative tolerance. The absolute drift is uniform across score magnitudes +# (~0.005-0.010), so it only exceeds the relative tolerance for the small +# ~0.10 text-vs-text value. Keep the relative tolerances tight and add only a +# small absolute floor for the affected backends. +# TRITON_ATTN: gfx942/ROCm 7.2 drifts ~0.008 abs on text-vs-text (~7.9% rel). BACKEND_ABS_TOL: dict[str, float] = { "default": 0.0, "ROCM_AITER_FA": 0.005, + "TRITON_ATTN": 0.009, "FLEX_ATTENTION": 0.006, } From 9fd00ee006ccd4996bbc756397b039343d2fde94 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 27 Jun 2026 04:08:54 -0500 Subject: [PATCH 071/138] [ROCm][CI] Move remaining mi250_2 tests out of the MI250 queue (#46905) Signed-off-by: Codex Co-authored-by: Codex --- .buildkite/test-amd.yaml | 214 ++++++++++++--------------------------- 1 file changed, 64 insertions(+), 150 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 083aa024b2a..a9608cc332f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -112,33 +112,6 @@ steps: # # ######################################################################################################################################### -#----------------------------------------------------- mi250 ยท basic_correctness -----------------------------------------------------# - -- 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=MI250 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 ยท compile ----------------------------------------------------------# - label: PyTorch Compilation Unit Tests # TBD @@ -179,48 +152,8 @@ steps: commands: - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;" -- 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: - - 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: Distributed Comm Ops # 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/distributed - - tests/distributed - - 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 - - label: Pipeline + Context Parallelism (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] @@ -330,54 +263,6 @@ 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: 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 - - 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: `colbert_query io_processor` test - - pip install -e ./plugins/colbert_query_plugin - - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py - - pip uninstall colbert_query_plugin -y - # END: `colbert_query 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 plugins_tests/test_oot_registration_online.py # it needs a clean process - - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process - - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins - #------------------------------------------------------------ mi250 ยท v1 -------------------------------------------------------------# - label: Batch Invariance (H100-MI250) # TBD @@ -505,41 +390,6 @@ steps: commands: - pytest -v -s v1/attention -- label: Distributed DP 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/v1/distributed - - tests/entrypoints/openai/test_multi_api_servers.py - - vllm/platforms/rocm.py - commands: - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - - 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: 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" - #------------------------------------------------------------- mi250 ยท misc ------------------------------------------------------------# - label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD @@ -782,6 +632,22 @@ steps: #-------------------------------------------------------- mi300 ยท distributed --------------------------------------------------------# +- label: Distributed Comm Ops # 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/distributed + - tests/distributed + - 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 + - label: EPLB Algorithm # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1806,6 +1672,54 @@ steps: - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper +#---------------------------------------------------------- mi300 ยท plugins ----------------------------------------------------------# + +- label: Plugin 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/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 + - 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: `colbert_query io_processor` test + - pip install -e ./plugins/colbert_query_plugin + - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py + - pip uninstall colbert_query_plugin -y + # END: `colbert_query 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 plugins_tests/test_oot_registration_online.py # it needs a clean process + - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process + - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins + #------------------------------------------------------- mi300 ยท quantization --------------------------------------------------------# - label: Quantization # TBD From 867fd5e8ed6b0bfcf84b24f82658c9fb698a6d35 Mon Sep 17 00:00:00 2001 From: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Date: Sat, 27 Jun 2026 06:22:57 -0400 Subject: [PATCH 072/138] [ROCm][Perf] Use flydsl moe with Minimax-M3 mxfp8 weights on gfx950 and implemented moe-backend selection (#46184) Signed-off-by: Hongxia Yang Signed-off-by: tjtanaa Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: TJian Co-authored-by: Andreas Karatzas Co-authored-by: Tan Pin Siang --- .../moe/test_mxfp8_aiter_backend_selection.py | 135 ++++++++++++++ vllm/_aiter_ops.py | 30 ++++ .../fused_moe/experts/aiter_mxfp8_moe.py | 166 ++++++++++++++++++ .../layers/fused_moe/oracle/fp8.py | 10 +- .../layers/fused_moe/oracle/mxfp8.py | 37 +++- 5 files changed, 374 insertions(+), 4 deletions(-) create mode 100644 tests/kernels/moe/test_mxfp8_aiter_backend_selection.py create mode 100644 vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py new file mode 100644 index 00000000000..7c2fdbabe29 --- /dev/null +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 MoE backend selection for the AITER FlyDSL kernel (gfx950). + +GPU-free: mocks the platform (gfx950) and the ``flydsl`` package check, then +exercises the oracle so the FlyDSL backend is auto-picked when usable (including +under expert parallelism, since apply() forwards the expert_map as aiter's +expert_mask) and skipped (native fallback) when the device/package is missing. +""" + +import dataclasses +from unittest.mock import patch + +import pytest + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("This test can only run on ROCm.", allow_module_level=True) + +from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402 +from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402 + AiterMxfp8Experts, +) +from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402 + FusedMoEActivationFormat, +) +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( # noqa: E402 + Fp8MoeBackend, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402 + _BACKEND_NAME_MAP, + _SUPPORTED_BACKENDS, + _mxfp8_backend_to_kernel_cls, + _select_kernel_cls, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 + kMxfp8Dynamic, + kMxfp8Static, +) + +_AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe" + + +def _config(ep_size: int = 1): + cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144) + if ep_size != 1: + cfg = dataclasses.replace( + cfg, + moe_parallel_config=dataclasses.replace( + cfg.moe_parallel_config, ep_size=ep_size, use_ep=True + ), + ) + return cfg + + +def _gfx950(): + """Patch the platform so the device gate (gfx950 / MX) passes off-ROCm.""" + return patch.multiple( + f"{_AITER_MOD}.current_platform", + is_rocm=lambda: True, + supports_mx=lambda: True, + ) + + +def _flydsl_installed(present: bool): + return patch(f"{_AITER_MOD}.is_aiter_mxfp8_moe_available", return_value=present) + + +def test_aiter_mxfp8_registered(): + """The FlyDSL backend is auto-selectable and reachable via --moe-backend aiter.""" + assert Fp8MoeBackend.AITER_MXFP8 in _SUPPORTED_BACKENDS + assert _BACKEND_NAME_MAP["aiter"] is Fp8MoeBackend.AITER_MXFP8 + assert _mxfp8_backend_to_kernel_cls(Fp8MoeBackend.AITER_MXFP8) == [ + AiterMxfp8Experts + ] + + +def test_triton_selectable(): + assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8 + # Not auto-selected (only reachable explicitly), so FlyDSL still wins auto. + assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS + + +@pytest.mark.parametrize("ep_size", [1, 2]) +def test_ep_supported(ep_size): + """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask.""" + assert ( + AiterMxfp8Experts._supports_parallel_config( + _config(ep_size).moe_parallel_config + ) + is True + ) + + +@pytest.mark.parametrize( + "present,ep_size,supported,reason_substr", + [ + (True, 1, True, None), # gfx950 + flydsl + TP -> selectable + (True, 2, True, None), # gfx950 + flydsl + EP -> selectable (expert_mask) + (False, 1, False, "flydsl package"), # package missing -> clear reason + ], +) +def test_is_supported_config(present, ep_size, supported, reason_substr): + with _gfx950(), _flydsl_installed(present): + ok, reason = AiterMxfp8Experts.is_supported_config( + AiterMxfp8Experts, + _config(ep_size), + kMxfp8Static, + kMxfp8Dynamic, + FusedMoEActivationFormat.Standard, + ) + assert ok is supported + if reason_substr is not None: + assert reason_substr in reason + + +def test_explicit_moe_backend_aiter(): + """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear + ValueError when the flydsl package is missing.""" + with _gfx950(), _flydsl_installed(True): + assert ( + _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) + is AiterMxfp8Experts + ) + assert ( + _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(2)) + is AiterMxfp8Experts + ) + with ( + _gfx950(), + _flydsl_installed(False), + pytest.raises(ValueError, match="flydsl package"), + ): + _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 95a5361032f..4a8b4209d87 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -171,6 +171,7 @@ def _rocm_aiter_fused_moe_impl( bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, + swiglu_limit: float = 0.0, ) -> torch.Tensor: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe @@ -203,6 +204,7 @@ def _rocm_aiter_fused_moe_impl( bias1=bias1, bias2=bias2, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, + swiglu_limit=swiglu_limit, **extra_kwargs, ) @@ -229,6 +231,7 @@ def _rocm_aiter_fused_moe_fake( bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, + swiglu_limit: float = 0.0, ) -> torch.Tensor: if output_dtype is not None: return torch.empty_like(hidden_states, dtype=output_dtype) @@ -2201,6 +2204,7 @@ class rocm_aiter_ops: bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, + swiglu_limit: float = 0.0, ) -> torch.Tensor: return torch.ops.vllm.rocm_aiter_fused_moe( hidden_states, @@ -2224,6 +2228,7 @@ class rocm_aiter_ops: bias1, bias2, moe_sorting_dispatch_policy, + swiglu_limit, ) @staticmethod @@ -2678,6 +2683,31 @@ class rocm_aiter_ops: return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors) + @staticmethod + def shuffle_mxfp8_moe_weights( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Preshuffle MXFP8 MoE weights + E8M0 scales into AITER's FlyDSL layout: + gate/up-interleaved weights, interleaved scale for w13 (gate/up), plain + scale for w2 (the interleaved variant is gate/up-only and misaligns w2). + """ + from aiter.ops.shuffle import shuffle_scale, shuffle_weight + + num_experts = w13.shape[0] + w13 = shuffle_weight(w13, is_guinterleave=True, gate_up=True) + w2 = shuffle_weight(w2, is_guinterleave=True, gate_up=False) + w13_scale = shuffle_scale( + w13_scale.reshape(-1, w13_scale.shape[-1]), + num_experts, + is_guinterleave=True, + gate_up=True, + ) + w2_scale = shuffle_scale(w2_scale.reshape(-1, w2_scale.shape[-1])) + return w13, w2, w13_scale, w2_scale + @staticmethod def flash_attn_varlen_func( q: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py new file mode 100644 index 00000000000..3cbab0a0d54 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 (1x32 block, E8M0) MoE via AITER's FlyDSL two-stage grouped GEMM +(gfx950); alternative to ``Mxfp8NativeTritonExperts``. Routes through +``aiter.fused_moe`` (per_1x32, gate_mode=INTERLEAVE); weights are preshuffled in +``convert_to_fp8_moe_kernel_format``. +""" + +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.experts.mxfp8_emulation_moe import ( + Mxfp8TritonExpertsBase, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def is_aiter_mxfp8_moe_available() -> bool: + """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl`` + package is importable, AND the installed aiter carries the mxfp8 FlyDSL + 2-stage support from ROCm/aiter#3811. + + ``flydsl`` and ``aiter`` are separate packages, so ``is_flydsl_available()`` + (flydsl pkg + arch) is necessary but not sufficient: an older aiter without + #3811 still ships the flydsl pkg and the ``aiter.ops.flydsl`` module but a + broken/missing ``per_1x32 + fp8`` 2-stage path. Without this extra gate a + nightly lacking #3811 would wrongly select FlyDSL instead of falling back to + the native Triton dot_scaled path. #3811 added no probe-able public symbol, + so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails + closed (returns False -> triton dot_scaled), which is always safe.""" + if not (current_platform.is_rocm() and current_platform.supports_mx()): + return False + try: + import os + + import aiter + from aiter.ops.flydsl.utils import is_flydsl_available + + if not is_flydsl_available(): + return False + return os.path.exists( + os.path.join( + os.path.dirname(aiter.__file__), + "configs", + "model_configs", + "minimax_m3_mxfp8_tuned_fmoe.csv", + ) + ) + except Exception: + return False + + +class AiterMxfp8Experts(Mxfp8TritonExpertsBase): + """MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950).""" + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return self.quant_config.quant_dtype + + @property + def block_shape(self) -> list[int] | None: + return self.quant_config.block_shape + + @property + def expects_unquantized_inputs(self) -> bool: + # aiter.fused_moe MXFP8-quantizes the activations internally. + return True + + @staticmethod + def _supports_current_device() -> bool: + # Device capability only (gfx950 / MX-capable ROCm). The flydsl package + # check lives in is_supported_config so a missing package is reported + # distinctly from an unsupported device. + return current_platform.is_rocm() and current_platform.supports_mx() + + @staticmethod + def _supports_parallel_config(moe_parallel_config) -> bool: + # Both TP (expert_map=None) and EP are supported: apply() forwards the + # expert_map as aiter's ``expert_mask`` (the per-rank local-expert + # selection), mirroring the native rocm_aiter_moe path. + return True + + @staticmethod + def is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ): + is_supported, reason = super().is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + # _supports_current_device() only gates on the device; surface a clear + # reason when the device is fine but the flydsl package is missing. + if is_supported and not is_aiter_mxfp8_moe_available(): + return False, ( + "kernel requires the aiter flydsl package, which is not installed" + ) + return is_supported, reason + + 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, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + from aiter import ActivationType, QuantType + from aiter.ops.flydsl.moe_common import GateMode + + from vllm._aiter_ops import rocm_aiter_ops + + # Re-tag the preshuffled weights: replace_parameter drops the + # is_shuffled flag, without which aiter picks a broken CK kernel. + w1.is_shuffled = True + w2.is_shuffled = True + + limit = self.quant_config.gemm1_clamp_limit + swiglu_limit = 0.0 if limit is None else float(limit) + + # Under EP, aiter expects ``expert_mask`` as a 0/1 *local-expert* mask + # over global ids with a trailing fake-expert sentinel slot + # (shape ``[global_num_experts + 1]``), NOT vLLM's expert_map (a + # global->local index map with -1 for non-local). Convert it; aiter + # derives the global->local compaction from the mask itself. ``None`` + # under pure TP. + if expert_map is not None: + local_mask = (expert_map >= 0).to(torch.int32) + expert_mask = torch.cat([local_mask, local_mask.new_zeros(1)]) + else: + expert_mask = None + + # Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the + # call is captured under HIP graphs / torch.compile (a direct + # ``aiter.fused_moe`` is opaque to the dispatcher). aiter requires FP32 + # routing weights / INT32 ids. + out = rocm_aiter_ops.fused_moe( + hidden_states, + w1, + w2, + topk_weights.to(torch.float32), + topk_ids.to(torch.int32), + expert_mask=expert_mask, + activation_method=ActivationType.Swiglu.value, + quant_method=QuantType.per_1x32.value, + doweight_stage1=apply_router_weight_on_input, + w1_scale=self.w1_scale_val, + w2_scale=self.w2_scale_val, + a1_scale=None, + a2_scale=None, + gate_mode=GateMode.INTERLEAVE.value, + swiglu_limit=swiglu_limit, + output_dtype=output.dtype, + ) + output.copy_(out.to(output.dtype)) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 9f930f1d58c..862f292009c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -59,7 +59,9 @@ class Fp8MoeBackend(Enum): # MXFP8 MoE via a Triton ``dot_scaled`` kernel that lowers to CDNA4 # (gfx950) native MX matrix-core ops. Weights stay in MXFP8 (no load-time # format conversion); the FP8 values + E8M0 scales are consumed directly. - NATIVE_MXFP8 = "NATIVE_MXFP8" + TRITON_MXFP8 = "TRITON_MXFP8" + # MXFP8 MoE via AITER (FlyDSL two-stage grouped GEMM) on gfx950. + AITER_MXFP8 = "AITER_MXFP8" def _get_priority_backends( @@ -423,6 +425,10 @@ def convert_to_fp8_moe_kernel_format( ) elif fp8_backend == Fp8MoeBackend.AITER: w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2) + elif fp8_backend == Fp8MoeBackend.AITER_MXFP8: + w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights( + w13, w2, w13_scale, w2_scale + ) elif fp8_backend == Fp8MoeBackend.MARLIN: weight_block_size = getattr(layer, "weight_block_size", None) if weight_block_size == [1, 32]: @@ -484,7 +490,7 @@ def convert_to_fp8_moe_kernel_format( # EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes # the MXFP8 weights as-is โ€” neither needs a load-time layout change. Fp8MoeBackend.EMULATION, - Fp8MoeBackend.NATIVE_MXFP8, + Fp8MoeBackend.TRITON_MXFP8, ]: raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}") diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index d0d7c76481b..06b622a6c4b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -21,6 +21,10 @@ _SUPPORTED_BACKENDS = ( Fp8MoeBackend.DEEPGEMM, Fp8MoeBackend.MARLIN, Fp8MoeBackend.XPU, + # AITER FlyDSL (gfx950): auto-picked by select_mxfp8_moe_backend when + # is_supported_config passes (gfx950 + flydsl installed + not EP). On other + # devices / no flydsl / EP it is skipped and native is used. + Fp8MoeBackend.AITER_MXFP8, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -28,6 +32,8 @@ _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { "deep_gemm": Fp8MoeBackend.DEEPGEMM, "marlin": Fp8MoeBackend.MARLIN, "xpu": Fp8MoeBackend.XPU, + "aiter": Fp8MoeBackend.AITER_MXFP8, + "triton": Fp8MoeBackend.TRITON_MXFP8, } @@ -46,6 +52,27 @@ def _mxfp8_backend_to_kernel_cls( ) return [DeepGemmExperts] + if backend == Fp8MoeBackend.AITER_MXFP8: + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( + AiterMxfp8Experts, + ) + + return [AiterMxfp8Experts] + if backend == Fp8MoeBackend.TRITON_MXFP8: + # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e. + # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise. + # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``. + if current_platform.supports_mx(): + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( + Mxfp8NativeTritonExperts, + ) + + return [Mxfp8NativeTritonExperts] + from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( + Mxfp8EmulationTritonExperts, + ) + + return [Mxfp8EmulationTritonExperts] return backend_to_kernel_cls(backend) @@ -77,7 +104,13 @@ def _select_kernel_cls( def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: - """ROCm fallback when vendor MXFP8 backends are unavailable.""" + """ROCm fallback when no auto-selected MXFP8 backend is available. + + The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by + ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or + explicitly via ``--moe-backend aiter``; this fallback handles the rest + (native dot_scaled on gfx950, else BF16 emulation). + """ if current_platform.supports_mx(): from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( @@ -85,7 +118,7 @@ def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts ) logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.") - return Fp8MoeBackend.NATIVE_MXFP8, Mxfp8NativeTritonExperts + return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8EmulationTritonExperts, From 51a99565c398c8320de8131e07731c75c52eb87c Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sat, 27 Jun 2026 05:34:17 -0700 Subject: [PATCH 073/138] [ROCm][Perf] Fused shared expert for Minimax M3 (#46474) Signed-off-by: Fangzhou-Ai Signed-off-by: tjtanaa Co-authored-by: Claude Opus 4.8 Co-authored-by: tjtanaa --- vllm/models/minimax_m3/amd/model.py | 59 ++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 13 deletions(-) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 894550c1576..7bb8bd722f2 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -23,8 +23,9 @@ import torch from torch import nn from transformers import PretrainedConfig -import vllm.envs as envs from vllm import _custom_ops as ops +from vllm import envs +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import ( CacheConfig, @@ -96,7 +97,6 @@ from vllm.models.minimax_m3.common.sparse_attention import ( ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype from vllm.v1.kv_cache_interface import ( @@ -107,14 +107,15 @@ from vllm.v1.kv_cache_interface import ( def _fuse_shared_experts_enabled(config: PretrainedConfig) -> bool: - """Whether to fuse the shared expert into the routed grouped MoE. + """Whether to fuse the shared expert with routed experts. ROCm only. Opt-in via ``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS`` (the - router-append fusion runs on the triton/flydsl mxfp8 MoE independent of the - aiter master switch); requires a shared expert and is disabled under expert - parallelism (the shared slot is appended to the routed top-k, which the EP - expert-map path does not handle). + router-append fusion runs on both aiter and non-aiter MoE); + it is disabled under expert parallelism (the shared slot is appended to + the routed top-k, which the EP expert-mapping path does not handle). """ + from vllm.platforms import current_platform + return bool( current_platform.is_rocm() and getattr(config, "n_shared_experts", None) @@ -268,6 +269,24 @@ class MiniMaxM3MLP(nn.Module): return x +def _aiter_moe_fused_shared_experts_enabled(config: PretrainedConfig) -> bool: + """Whether the fused shared expert routes through aiter's grouped top-k MoE. + + A strict sub-case of :func:`_fuse_shared_experts_enabled`: shared-expert + fusion must already be opted in (``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS``) + and allowed (not under expert parallelism). When additionally on gfx950 with + an active aiter MoE backend, the shared expert is appended inside aiter's + biased grouped top-k kernel (``num_fused_shared_experts``) instead of the + vLLM router's torch concat. Otherwise FSE still runs via the vLLM top-k bias + router. + """ + if not _fuse_shared_experts_enabled(config): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + + class MiniMaxM3MoE(nn.Module): """Sigmoid-routed MoE block with a routing-bias correction and a shared expert.""" @@ -313,11 +332,14 @@ class MiniMaxM3MoE(nn.Module): prefix=f"{prefix}.gate", ) - # Fuse the shared expert into the routed grouped GEMM when opted in via - # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: it becomes routed-expert slot - # ``num_local_experts``, reached by every token, eliminating the - # separate dense-MLP launches. Not supported under expert parallelism. + # Shared-expert fusion (opt-in via VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS, + # off under expert parallelism) folds the shared expert into the routed + # MoE call as the last expert slot, so we don't build a separate module. + # On gfx950 with aiter MoE the append is fused inside aiter's grouped + # top-k kernel; otherwise it goes through the vLLM top-k bias router. self.fuse_shared_experts = _fuse_shared_experts_enabled(config) + self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled(config) + self.shared_experts: MiniMaxM3MLP | None = None if self.n_shared_experts and not self.fuse_shared_experts: self.shared_experts = MiniMaxM3MLP( @@ -328,6 +350,13 @@ class MiniMaxM3MoE(nn.Module): prefix=f"{prefix}.shared_experts", ) + # The aiter MoE fused path goes through aiter's biased grouped top-k + # (GroupedTopKRouter, as in DeepSeek-V4): M3 is not group-routed, so a + # trivial single group (num_expert_group=topk_group=1) reduces to plain + # top-k while applying the sigmoid + bias correction and appending the + # always-on shared expert; aiter applies the routed scaling internally. + # Every other path (vLLM top-k bias router, or no fusion) applies the + # routed scaling to the MoE output here. self.experts = FusedMoE( num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, @@ -337,12 +366,15 @@ class MiniMaxM3MoE(nn.Module): scoring_func=config.scoring_func, e_score_correction_bias=self.e_score_correction_bias, renormalize=True, + use_grouped_topk=self.use_aiter_moe_fse, + num_expert_group=1 if self.use_aiter_moe_fse else None, + topk_group=1 if self.use_aiter_moe_fse else None, activation="swigluoai_uninterleave", swiglu_limit=config.swiglu_limit, swiglu_alpha=config.swiglu_alpha, swiglu_beta=config.swiglu_beta, routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_output=True, + apply_routed_scale_to_output=not self.use_aiter_moe_fse, router_logits_dtype=self.gate.out_dtype, shared_experts=self.shared_experts, n_shared_experts=( @@ -927,9 +959,10 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = self.get_expert_mapping() + _fuse_shared = _fuse_shared_experts_enabled(self.config) + params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - _fuse_shared = _fuse_shared_experts_enabled(self.config) for name, loaded_weight in weights: # The MTP module is not modeled yet. if "mtp." in name: From 35e3850fa9499b99f0b32ee8e9d5551a290d9c54 Mon Sep 17 00:00:00 2001 From: xiaolinchen <2990624738@qq.com> Date: Sun, 28 Jun 2026 02:30:10 +0800 Subject: [PATCH 074/138] [Bugfix][Test] Fix test_flashinfer_cutlass_mxfp4_fused_moe on sm90 (stale weight/scale interleave) (#46915) Signed-off-by: wentian-byte <2990624738@qq.com> --- tests/kernels/moe/test_ocp_mx_moe.py | 29 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index a96e47fe439..8c620afbc81 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -659,19 +659,6 @@ def test_trtllm_gen_mxfp4_fused_moe( check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8) -def _interleave_scales_lastdim_by4(scales: torch.Tensor) -> torch.Tensor: - """Interleave scales on the last dimension by groups of 4, matching - the transformation in mxfp4.py's BF16 (Hopper) path.""" - s = scales.to(torch.uint8) - s_shape = s.shape - assert s_shape[-1] % 4 == 0 - s = s.reshape(*s_shape[:-1], s_shape[-1] // 4, 4) - # Move the 4-group dimension before the row dimension - permuted = s.permute(0, 2, 1, 3) - # Merge the row dim with the 4-group dim - return permuted.reshape(s_shape[0], s_shape[-1] // 4, s_shape[1] * 4) - - @pytest.mark.parametrize("topk", [1, 4]) @pytest.mark.parametrize("num_experts", [32]) @pytest.mark.parametrize("num_tokens", [1, 128]) @@ -771,13 +758,25 @@ def test_flashinfer_cutlass_mxfp4_fused_moe( w1_w, w3_w = torch.chunk(w13_q, 2, dim=1) w13_q_swapped = torch.cat([w3_w, w1_w], dim=1) + # SM90 mixed-input GEMM expects weights/scales in an interleaved layout; + # without it the FP4->BF16 LUT reads bytes from wrong positions for K>128. + from flashinfer.fused_moe import ( + interleave_moe_scales_for_sm90_mixed_gemm, + interleave_moe_weights_for_sm90_mixed_gemm, + ) + + w13_q_swapped = interleave_moe_weights_for_sm90_mixed_gemm( + w13_q_swapped, quant_type="fp4" + ) + w2_q = interleave_moe_weights_for_sm90_mixed_gemm(w2_q, quant_type="fp4") + b1, b3 = torch.chunk(bias13.to(torch.float32), 2, dim=-1) w13_b = torch.cat([b3, b1], dim=-1).to(torch.bfloat16) w1_s, w3_s = torch.chunk(w13_scale, 2, dim=1) w13_s = torch.cat([w3_s, w1_s], dim=1) - w13_s_inter = _interleave_scales_lastdim_by4(w13_s) - w2_s_inter = _interleave_scales_lastdim_by4(w2_scale) + w13_s_inter = interleave_moe_scales_for_sm90_mixed_gemm(w13_s) + w2_s_inter = interleave_moe_scales_for_sm90_mixed_gemm(w2_scale) routing_weights = torch.nn.functional.softmax( router_logits, dim=1, dtype=torch.float32 From 56aa067bf05a7bc26f0fa017774e8521ccae7144 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 27 Jun 2026 16:17:33 -0400 Subject: [PATCH 075/138] [CI Bug] Fix h100 `AssertionError: Cold-start child failed` (#46927) Signed-off-by: yewentao256 --- tests/compile/h100/test_startup.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/compile/h100/test_startup.py b/tests/compile/h100/test_startup.py index 78554a3e93d..075fc8e2497 100644 --- a/tests/compile/h100/test_startup.py +++ b/tests/compile/h100/test_startup.py @@ -138,10 +138,10 @@ MODEL_SPECS = [ ModelStartupSpec( model="deepseek-ai/DeepSeek-V3.2", hf_overrides=_SMALL_MOE_OVERRIDES, - cold_artifacts_saved=4, + cold_artifacts_saved=9, # https://github.com/vllm-project/vllm/issues/38051 - warm_artifacts_saved=0 if is_torch_equal_or_newer("2.12.0") else 4, - warm_artifacts_loaded=4 if is_torch_equal_or_newer("2.12.0") else 0, + warm_artifacts_saved=0 if is_torch_equal_or_newer("2.12.0") else 9, + warm_artifacts_loaded=9 if is_torch_equal_or_newer("2.12.0") else 0, ), id="deepseek_v3.2", ), From ea2ead1db33dafb067aa64d4ce7b9c2150c12091 Mon Sep 17 00:00:00 2001 From: jj shao Date: Sun, 28 Jun 2026 04:23:59 +0800 Subject: [PATCH 076/138] [Misc] Fix incorrect layer type annotation in Fp8LinearMethod (#46818) Signed-off-by: shaojinjie.sjj Co-authored-by: shaojinjie.sjj --- vllm/model_executor/layers/quantization/fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 869fbf75237..d4ec1c093a8 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -320,7 +320,7 @@ class Fp8LinearMethod(LinearMethodBase): def create_weights( self, - layer: RoutedExperts, + layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, @@ -394,7 +394,7 @@ class Fp8LinearMethod(LinearMethodBase): self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.use_marlin: if not self.block_quant: # Canonicalize to (K, N) for the kernel. From 8bf064f8d3408ca89cabc2f071adc696314c867e Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sat, 27 Jun 2026 16:57:47 -0400 Subject: [PATCH 077/138] Fixed chunked embedding aggregation with request-id metadata (#46782) Signed-off-by: Taneem Ibrahim --- .../pooling/embed/test_io_processor.py | 92 ++++++++++++++ .../entrypoints/pooling/embed/io_processor.py | 120 +++++++++--------- vllm/entrypoints/pooling/typing.py | 7 + 3 files changed, 157 insertions(+), 62 deletions(-) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index fbee91fc48e..5a7a8aab2a6 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -3,6 +3,7 @@ """Unit tests for EmbedIOProcessor.""" import pytest +import torch from pydantic import TypeAdapter, ValidationError from vllm import PoolingParams @@ -19,6 +20,7 @@ from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingRequest, ) from vllm.entrypoints.pooling.typing import PoolingServeContext +from vllm.outputs import PoolingOutput, PoolingRequestOutput class TestEmbeddingRequestParsing: @@ -398,6 +400,96 @@ class TestValidateInputType: handler._validate_input_type("z") +class TestChunkedEmbeddingProcessing: + """Unit tests for chunked embedding aggregation.""" + + class _FakeModelConfig: + max_model_len = 3 + + @classmethod + def _make_handler(cls): + handler = object.__new__(EmbedIOProcessor) + handler.model_config = cls._FakeModelConfig() + return handler + + @staticmethod + def _make_context() -> PoolingServeContext[EmbeddingCompletionRequest]: + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "input": [[0, 1, 2, 3, 4], [10, 11]], + } + ) + assert isinstance(request, EmbeddingCompletionRequest) + return PoolingServeContext( + request=request, + pooling_params=PoolingParams(), + model_name="test", + request_id="embd-client-prompt-999-chunk-888", + engine_inputs=[ + {"prompt_token_ids": [0, 1, 2, 3, 4]}, + {"prompt_token_ids": [10, 11]}, + ], + ) + + @staticmethod + def _make_output( + request_id: str, + prompt_token_ids: list[int], + embedding: list[float], + ) -> PoolingRequestOutput: + return PoolingRequestOutput( + request_id=request_id, + outputs=PoolingOutput(data=torch.tensor(embedding)), + prompt_token_ids=prompt_token_ids, + num_cached_tokens=0, + finished=True, + ) + + def test_aggregation_uses_metadata_not_request_id_parsing(self): + handler = self._make_handler() + ctx = self._make_context() + + handler._pre_process_chunked(ctx) + + assert ctx.prompt_request_ids == [ + "embd-client-prompt-999-chunk-888-prompt-0-chunk-0", + "embd-client-prompt-999-chunk-888-prompt-0-chunk-1", + "embd-client-prompt-999-chunk-888-prompt-1-chunk-0", + ] + assert ctx.chunked_embedding_metadata is not None + assert [ + (item.prompt_index, item.chunk_index) + for item in ctx.chunked_embedding_metadata + ] == [(0, 0), (0, 1), (1, 0)] + + ctx.final_res_batch = [ + self._make_output(ctx.prompt_request_ids[0], [0, 1, 2], [1.0, 1.0]), + self._make_output(ctx.prompt_request_ids[1], [3, 4], [4.0, 7.0]), + self._make_output(ctx.prompt_request_ids[2], [10, 11], [9.0, 9.0]), + ] + + handler._post_process_chunked(ctx) + + assert len(ctx.final_res_batch) == 2 + assert ctx.final_res_batch[0].request_id == ( + "embd-client-prompt-999-chunk-888-prompt-0" + ) + assert ctx.final_res_batch[0].prompt_token_ids == [0, 1, 2, 3, 4] + assert torch.allclose( + ctx.final_res_batch[0].outputs.data, + torch.tensor([2.2, 3.4]), + ) + assert ctx.final_res_batch[1].request_id == ( + "embd-client-prompt-999-chunk-888-prompt-1" + ) + assert ctx.final_res_batch[1].prompt_token_ids == [10, 11] + assert torch.allclose( + ctx.final_res_batch[1].outputs.data, + torch.tensor([9.0, 9.0]), + ) + + class TestPreProcessCohereOnline: """Unit tests for EmbedIOProcessor._pre_process_cohere_online.""" diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index d2e6f23c149..ec52e2efd68 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from dataclasses import dataclass from typing import Any, Literal, cast import torch @@ -27,6 +28,7 @@ from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( + ChunkedEmbeddingMetadata, OfflineInputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, @@ -46,6 +48,12 @@ from .protocol import ( logger = init_logger(__name__) +@dataclass +class _ChunkedPromptAggregator: + weighted_sum: torch.Tensor | None = None + total_weight: int = 0 + + class EmbedIOProcessor(PoolingIOProcessor): name = "embed" @@ -113,6 +121,7 @@ class EmbedIOProcessor(PoolingIOProcessor): max_model_len = self.model_config.max_model_len chunked_engine_inputs: list[EngineInput] = [] prompt_request_ids: list[str] = [] + chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] = [] for prompt_idx, engine_input in enumerate(ctx.engine_inputs): token_ids = engine_input.get("prompt_token_ids", None) if token_ids is None: @@ -132,9 +141,16 @@ class EmbedIOProcessor(PoolingIOProcessor): prompt_request_ids.append( f"{request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}" ) + chunked_embedding_metadata.append( + ChunkedEmbeddingMetadata( + prompt_index=prompt_idx, + chunk_index=chunk_idx, + ) + ) ctx.engine_inputs = chunked_engine_inputs ctx.prompt_request_ids = prompt_request_ids + ctx.chunked_embedding_metadata = chunked_embedding_metadata return None @@ -142,66 +158,48 @@ class EmbedIOProcessor(PoolingIOProcessor): # Online aggregation for chunked requests to # minimize memory usage # Track aggregation state for each prompt - prompt_aggregators: dict[int, dict[str, Any]] = {} - short_prompts_results: dict[int, PoolingRequestOutput] = {} - for result_idx, result in enumerate(ctx.final_res_batch): - if "-chunk-" not in result.request_id: - # Non-chunked result - extract prompt_idx from request_id - parts = result.request_id.split("-") - try: - # Last part should be prompt index - prompt_idx = int(parts[-1]) - except (ValueError, IndexError): - prompt_idx = result_idx # Fallback to result_idx + if ctx.chunked_embedding_metadata is None: + raise ValueError("Chunked embedding metadata not available") + if len(ctx.chunked_embedding_metadata) != len(ctx.final_res_batch): + raise ValueError( + "Chunked embedding metadata count does not match result count" + ) - short_prompts_results[prompt_idx] = result + prompt_aggregators: dict[int, _ChunkedPromptAggregator] = {} + for result, chunk_metadata in zip( + ctx.final_res_batch, ctx.chunked_embedding_metadata + ): + prompt_idx = chunk_metadata.prompt_index + aggregator = prompt_aggregators.setdefault( + prompt_idx, _ChunkedPromptAggregator() + ) + + # MEAN pooling with online weighted averaging + # Ensure result is PoolingRequestOutput + # for embedding processing + if not isinstance(result, PoolingRequestOutput): + raise ValueError( + f"Expected PoolingRequestOutput for " + f"chunked embedding, got " + f"{type(result).__name__}" + ) + if result.prompt_token_ids is None: + raise ValueError( + "prompt_token_ids cannot be None for chunked processing" + ) + + weight = len(result.prompt_token_ids) + embedding_data = result.outputs.data + weighted_embedding = embedding_data.to(dtype=torch.float32) * weight + + if aggregator.weighted_sum is None: + # First chunk + aggregator.weighted_sum = weighted_embedding else: - # Extract prompt_idx from chunked request_id - parts = result.request_id.split("-") - try: - prompt_idx = int(parts[parts.index("prompt") + 1]) - except (ValueError, IndexError): - # Fallback: extract from result_idx if parsing fails - prompt_idx = result_idx + # Accumulate + aggregator.weighted_sum += weighted_embedding - # Initialize aggregator for this prompt if needed - if prompt_idx not in prompt_aggregators: - prompt_aggregators[prompt_idx] = { - "weighted_sum": None, - "total_weight": 0, - "chunk_count": 0, - "request_id": result.request_id.split("-chunk-")[0], - } - - aggregator = prompt_aggregators[prompt_idx] - - # MEAN pooling with online weighted averaging - # Ensure result is PoolingRequestOutput - # for embedding processing - if not isinstance(result, PoolingRequestOutput): - raise ValueError( - f"Expected PoolingRequestOutput for " - f"chunked embedding, got " - f"{type(result).__name__}" - ) - if result.prompt_token_ids is None: - raise ValueError( - "prompt_token_ids cannot be None for chunked processing" - ) - - weight = len(result.prompt_token_ids) - embedding_data = result.outputs.data - weighted_embedding = embedding_data.to(dtype=torch.float32) * weight - - if aggregator["weighted_sum"] is None: - # First chunk - aggregator["weighted_sum"] = weighted_embedding - else: - # Accumulate - aggregator["weighted_sum"] += weighted_embedding - - aggregator["total_weight"] += weight - aggregator["chunk_count"] += 1 + aggregator.total_weight += weight if ctx.original_engine_inputs is None: raise ValueError("Original engine inputs not available") @@ -216,8 +214,8 @@ class EmbedIOProcessor(PoolingIOProcessor): # Finalize MEAN aggregation for this chunked prompt aggregator = prompt_aggregators[prompt_idx] - weighted_sum = aggregator["weighted_sum"] - total_weight = aggregator["total_weight"] + weighted_sum = aggregator.weighted_sum + total_weight = aggregator.total_weight if ( weighted_sum is not None @@ -243,7 +241,7 @@ class EmbedIOProcessor(PoolingIOProcessor): original_token_ids = cast(list[int], token_ids) pooling_request_output = PoolingRequestOutput( - request_id=aggregator["request_id"], + request_id=f"{ctx.request_id}-prompt-{prompt_idx}", prompt_token_ids=original_token_ids, outputs=pooling_output_data, num_cached_tokens=0, @@ -255,8 +253,6 @@ class EmbedIOProcessor(PoolingIOProcessor): raise ValueError( f"Failed to aggregate chunks for prompt {prompt_idx}" ) - elif prompt_idx in short_prompts_results: - final_res_batch.append(short_prompts_results[prompt_idx]) else: raise ValueError(f"Result not found for prompt {prompt_idx}") diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 2cf38490053..54d02c5b61f 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -70,6 +70,12 @@ AnyPoolingResponse: TypeAlias = ( PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest) +@dataclass(kw_only=True) +class ChunkedEmbeddingMetadata: + prompt_index: int + chunk_index: int + + @dataclass(kw_only=True) class PoolingServeContext(Generic[PoolingRequestT]): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -91,6 +97,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): ## for Long Text Embedding with Chunked Processing original_engine_inputs: Sequence[EngineInput] | None = None + chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] | None = None ## for bi-encoder & late-interaction n_queries: int | None = None From b6caeb5a0966103c6df22f019270d66233e1b687 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Sat, 27 Jun 2026 16:09:25 -0500 Subject: [PATCH 078/138] [Model Runner V2][Spec Decode] Use fp32 uniform threshold for acceptance (#46878) --- vllm/v1/worker/gpu/sample/gumbel.py | 11 +++++++++-- .../worker/gpu/spec_decode/rejection_sampler_utils.py | 6 +++--- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 6dbb04cd933..4b0a1694f70 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -73,6 +73,14 @@ def tl_rand64(seed, offset, includes_zero: tl.constexpr): return u +@triton.jit +def tl_rand32(seed, offset, includes_zero: tl.constexpr): + u = tl.rand(seed, offset) + if not includes_zero: + u = tl.maximum(u, _TL_RAND_MIN) + return u + + @triton.jit def gumbel_block_argmax( logits, @@ -131,8 +139,7 @@ def gumbel_block_argmax( u = tl_rand64(gumbel_seed, block, includes_zero=False) gumbel_noise = -tl.log(-tl.log(u)) else: - u = tl.rand(gumbel_seed, block) - u = tl.maximum(u, _TL_RAND_MIN) + u = tl_rand32(gumbel_seed, block, includes_zero=False) # Draw the large-noise tail (which decides the argmax winner) from u -> 0, # where fp32 has fine resolution, instead of u -> 1, where fp32 spacing is # ~2**-24. The naive `-log(-log(u))` puts the winning tail at u -> 1, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 7020f228046..070270c8324 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -3,7 +3,7 @@ import torch from vllm.triton_utils import tl, tldevice, triton -from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32 @triton.jit @@ -243,7 +243,7 @@ def _rejection_kernel( if SYNTHETIC_MODE: pos = tl.load(pos_ptr + logit_idx) - u = tl_rand64(seed, pos, includes_zero=False) + u = tl_rand32(seed, pos, includes_zero=False) rate = tl.load(synthetic_conditional_rates_ptr + i) # -1 is used for padded draft token ids that should be rejected. accepted &= (u < rate) & (draft_sampled >= 0) @@ -272,7 +272,7 @@ def _rejection_kernel( ) target_log_prob = target_logit - target_lse pos = tl.load(pos_ptr + logit_idx) - u = tl_rand64(seed, pos, includes_zero=False) + u = tl_rand32(seed, pos, includes_zero=False) if HAS_DRAFT_LOGITS: draft_logit = tl.load( draft_logits_ptr From 9036c89ee410b30913ca8b7d362a7d0805583b51 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sat, 27 Jun 2026 17:30:49 -0500 Subject: [PATCH 079/138] [Hardware][AMD][CI] Patch Whisper multi LoRA test to use TRITON_ATTN for now (#46928) Signed-off-by: Matthew Wong --- tests/lora/test_whisper.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/lora/test_whisper.py b/tests/lora/test_whisper.py index ea8179a9c66..6f1a894cf91 100644 --- a/tests/lora/test_whisper.py +++ b/tests/lora/test_whisper.py @@ -12,6 +12,7 @@ import pytest import vllm from vllm.assets.audio import AudioAsset from vllm.lora.request import LoRARequest +from vllm.platforms import current_platform from ..utils import create_new_process_for_each_test @@ -30,7 +31,9 @@ def use_spawn_for_whisper(monkeypatch): monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") -def create_whisper_llm(enable_lora: bool = True, max_loras: int = 2): +def create_whisper_llm( + enable_lora: bool = True, max_loras: int = 2, attn_backend: str | None = None +): """Create a Whisper LLM instance with optional LoRA support.""" return vllm.LLM( model=WHISPER_MODEL, @@ -40,6 +43,7 @@ def create_whisper_llm(enable_lora: bool = True, max_loras: int = 2): max_model_len=448, dtype="half", enforce_eager=True, # For stability in tests + attention_config={"backend": attn_backend}, ) @@ -109,7 +113,11 @@ def test_whisper_multi_lora(whisper_lora_files): This test verifies that the same LoRA adapter can be loaded with different IDs and produce consistent results. """ - llm = create_whisper_llm(enable_lora=True, max_loras=4) + llm = create_whisper_llm( + enable_lora=True, + max_loras=4, + attn_backend="TRITON_ATTN" if current_platform.is_rocm() else None, + ) # Test with different LoRA IDs using the same adapter outputs_lora1 = run_whisper_inference(llm, lora_path=whisper_lora_files, lora_id=1) From 798185d438c030f9b4fd62687440889e7b195251 Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Sat, 27 Jun 2026 21:01:45 -0400 Subject: [PATCH 080/138] [KV-Offloading] Fix tensors_per_block stride (#46888) Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath --- .../unit/offloading_connector/utils.py | 12 +++++++++++- .../kv_connector/v1/offloading/worker.py | 16 +++++++++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 482a2f25a56..c2884649bdd 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -41,6 +41,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheTensor, ) from vllm.v1.kv_offload.base import ( CanonicalKVCaches, @@ -241,9 +242,18 @@ class RequestRunner: ) ] + kv_cache_tensors = [ + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * num_gpu_blocks, + shared_by=[layer_name], + ) + for group in kv_cache_groups + for layer_name in group.layer_names + ] + kv_cache_config = KVCacheConfig( num_blocks=num_gpu_blocks, - kv_cache_tensors=[], + kv_cache_tensors=kv_cache_tensors, kv_cache_groups=kv_cache_groups, ) vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 254e0dec09f..29914e7388e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -53,6 +53,16 @@ class OffloadingConnectorWorker: kv_cache_config = self.spec.kv_cache_config num_blocks = kv_cache_config.num_blocks + # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use + # stride(0) as the manager-block stride (equals total_num_bytes_per_block). + # General (non-packed) layouts size the tensor at page_size_bytes per + # manager block, so page_size_bytes is the correct offloading stride. + layer_is_packed: dict[str, bool] = { + ln: bool(kv_tensor.block_stride) + for kv_tensor in kv_cache_config.kv_cache_tensors + for ln in kv_tensor.shared_by + } + # layer_name -> (num_blocks, page_size_bytes) tensor tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {} # layer_name -> size of (un-padded) page in bytes @@ -77,7 +87,11 @@ class OffloadingConnectorWorker: page = layer_kv_cache_spec.page_size_bytes elem_size = layer_kv_cache.element_size() byte_offset = layer_kv_cache.storage_offset() * elem_size - block_stride_bytes = layer_kv_cache.stride(0) * elem_size + block_stride_bytes = ( + layer_kv_cache.stride(0) * elem_size + if layer_is_packed[layer_name] + else page + ) tensors_per_block[layer_name] = ( torch.tensor( [], From 11a12305c0522c5c1ed273d7d3dc2304ac0cd495 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sun, 28 Jun 2026 09:38:07 +0800 Subject: [PATCH 081/138] [Model Runner V2][Spec Decode] Handle tuple hidden states from MTP draft models (#46786) --- .../test_gpu_autoregressive_speculator.py | 82 +++++++++++++++++++ .../spec_decode/autoregressive/speculator.py | 14 +--- .../gpu/spec_decode/gemma4/speculator.py | 7 -- .../worker/gpu/spec_decode/mtp/speculator.py | 4 - 4 files changed, 85 insertions(+), 22 deletions(-) create mode 100644 tests/v1/worker/test_gpu_autoregressive_speculator.py diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py new file mode 100644 index 00000000000..940fb375a92 --- /dev/null +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import nullcontext +from types import SimpleNamespace + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module +from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( + AutoRegressiveSpeculator, +) + + +class _TestSpeculator(AutoRegressiveSpeculator): + def load_draft_model(self, target_model, target_attn_layer_names): + raise NotImplementedError + + +class _DraftModel(torch.nn.Module): + def __init__(self, output: torch.Tensor | tuple[torch.Tensor, torch.Tensor]): + super().__init__() + self.output = output + + def forward(self, **kwargs): + return self.output + + +def _make_speculator( + monkeypatch, + output: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> _TestSpeculator: + monkeypatch.setattr( + spec_module, + "set_forward_context", + lambda *args, **kwargs: nullcontext(), + ) + + speculator = object.__new__(_TestSpeculator) + speculator.supports_mm_inputs = False + speculator.vllm_config = None + speculator.input_buffers = SimpleNamespace( + input_ids=torch.arange(4), + positions=torch.arange(4), + ) + speculator.hidden_states = torch.zeros(4, 3) + speculator.model = _DraftModel(output) + return speculator + + +def test_run_model_unpacks_tuple_return_for_mtp(monkeypatch): + logits_hidden = torch.full((4, 3), 1.0) + feedback_hidden = torch.full((4, 3), 2.0) + speculator = _make_speculator(monkeypatch, (logits_hidden, feedback_hidden)) + + actual_logits_hidden, actual_feedback_hidden = speculator._run_model( + 4, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + + assert actual_logits_hidden is logits_hidden + assert actual_feedback_hidden is feedback_hidden + + +def test_run_model_reuses_tensor_return_for_mtp(monkeypatch): + hidden = torch.full((4, 3), 1.0) + speculator = _make_speculator(monkeypatch, hidden) + + actual_logits_hidden, actual_feedback_hidden = speculator._run_model( + 4, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + + assert actual_logits_hidden is hidden + assert actual_feedback_hidden is hidden diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 775c06f7b8d..422d3ac6901 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -60,16 +60,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): """ return True - @property - def model_returns_tuple(self) -> bool: - """ - Whether the draft model's forward() returns a tuple. - - True: returns (last_hidden_states, hidden_states) โ€” Eagle, Gemma4 MTP. - False: returns a single tensor used for both โ€” standard MTP (DeepSeek). - """ - return True - def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: # Initialize cudagraph manager for draft prefill (draft position 0). self.prefill_cudagraph_manager = PrefillSpeculatorCudaGraphManager( @@ -328,7 +318,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): else: # Eager (NONE): call the raw model directly. ret_hidden_states = self.model(**model_inputs) - if self.model_returns_tuple: + # Some MTP models declare a single-tensor contract but return + # (logits_hidden, feedback_hidden) for final-norm correctness. + if isinstance(ret_hidden_states, tuple): last_hidden_states, hidden_states = ret_hidden_states else: last_hidden_states = ret_hidden_states diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index fcbea5d1012..dfa2c680109 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py @@ -30,13 +30,6 @@ class Gemma4Speculator(AutoRegressiveSpeculator): # No new KV slots are written, so positions and seq_lens stay fixed. return False - @property - def model_returns_tuple(self) -> bool: - # forward() returns (draft_hidden_states, backbone_hidden_states). - # The proposer uses draft_hidden_states for compute_logits and - # backbone_hidden_states for the hidden-state feedback buffer. - return True - def load_draft_model( self, target_model: nn.Module, diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index e6abb0be83a..4b9354f23e7 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -10,10 +10,6 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model class MTPSpeculator(AutoRegressiveSpeculator): - @property - def model_returns_tuple(self) -> bool: - return False - def load_draft_model( self, target_model: nn.Module, From a65f93fb2e295e501b929df3c291ec89c27d39e8 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 27 Jun 2026 23:51:19 -0500 Subject: [PATCH 082/138] [ROCm][CI] Add ci_base metadata for external cache orchestration (#46886) Signed-off-by: Andreas Karatzas Signed-off-by: Codex Co-authored-by: Codex --- .buildkite/scripts/ci-bake-rocm.sh | 194 ++++++++++++++++++++++++++--- docker/ci-rocm.hcl | 21 +++- 2 files changed, 194 insertions(+), 21 deletions(-) diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 1289939180d..51cffb8e20d 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -18,6 +18,7 @@ DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base tools/install_torchcodec_rocm.sh tests/vllm_test_utils" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base" +DEFAULT_CI_BASE_METADATA_VERSION="1" IMAGE_EXISTED_BEFORE_BUILD=0 TARGET="" @@ -525,6 +526,22 @@ get_remote_image_label_with_retry() { return 0 } +remote_ci_base_metadata_is_current() { + local image_ref="$1" + local metadata_version="" + + metadata_version=$(get_remote_image_label "${image_ref}" "vllm.ci_base.metadata_version") + [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]] +} + +remote_ci_base_metadata_is_current_with_retry() { + local image_ref="$1" + local metadata_version="" + + metadata_version=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.metadata_version") + [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]] +} + remote_image_exists() { local image_ref="$1" docker manifest inspect "${image_ref}" >/dev/null 2>&1 @@ -581,6 +598,7 @@ init_config() { CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" export PYTORCH_ROCM_ARCH @@ -635,6 +653,10 @@ load_ci_hcl() { echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}" } +init_bake_files() { + BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") +} + compute_ci_base_hash_if_needed() { if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then return 0 @@ -676,12 +698,14 @@ configure_ci_base_image_refs() { fi content_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${CI_BASE_CONTENT_HASH}") + CI_BASE_IMAGE_TAG_CONTENT_REF="${content_tag}" if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then commit_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${BUILDKITE_COMMIT}") - CI_BASE_IMAGE_TAG_COMMIT="${commit_tag}" - export CI_BASE_IMAGE_TAG_COMMIT fi + CI_BASE_IMAGE_TAG_COMMIT_REF="${commit_tag}" + # *_REF is the logical tag recorded in metadata. *_EXTRA is only passed to + # bake when that tag is not already the primary tag, avoiding duplicates. if should_push_stable_ci_base_tag; then primary_tag="${content_tag}" CI_BASE_IMAGE_TAG_STABLE="${stable_tag}" @@ -691,19 +715,33 @@ configure_ci_base_image_refs() { fi CI_BASE_IMAGE_TAG="${primary_tag}" if [[ "${primary_tag}" == "${content_tag}" ]]; then - CI_BASE_IMAGE_TAG_CONTENT="" + CI_BASE_IMAGE_TAG_CONTENT_EXTRA="" else - CI_BASE_IMAGE_TAG_CONTENT="${content_tag}" + CI_BASE_IMAGE_TAG_CONTENT_EXTRA="${content_tag}" fi - export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT CI_BASE_IMAGE_TAG_STABLE + if [[ -n "${commit_tag}" && "${commit_tag}" != "${primary_tag}" ]]; then + CI_BASE_IMAGE_TAG_COMMIT_EXTRA="${commit_tag}" + else + CI_BASE_IMAGE_TAG_COMMIT_EXTRA="" + fi + export CI_BASE_IMAGE_TAG + export CI_BASE_IMAGE_TAG_COMMIT_EXTRA + export CI_BASE_IMAGE_TAG_CONTENT_EXTRA + export CI_BASE_IMAGE_TAG_CONTENT_REF + export CI_BASE_IMAGE_TAG_COMMIT_REF + export CI_BASE_IMAGE_TAG_STABLE if is_ci_base_target; then IMAGE_TAG="${primary_tag}" export IMAGE_TAG echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}" - if [[ -n "${CI_BASE_IMAGE_TAG_COMMIT:-}" ]]; then - echo "ci_base commit image tag: ${CI_BASE_IMAGE_TAG_COMMIT}" + if [[ -n "${commit_tag}" ]]; then + if [[ "${commit_tag}" == "${primary_tag}" ]]; then + echo "ci_base commit image tag: ${commit_tag} (primary)" + else + echo "ci_base commit image tag: ${commit_tag}" + fi fi echo "ci_base content image tag: ${content_tag}" if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then @@ -728,8 +766,8 @@ ci_base_candidate_refs() { printf '%s\n' \ "${IMAGE_TAG:-}" \ "${CI_BASE_IMAGE_TAG:-}" \ - "${CI_BASE_IMAGE_TAG_COMMIT:-}" \ - "${CI_BASE_IMAGE_TAG_CONTENT:-}" \ + "${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}" \ + "${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}" \ "${CI_BASE_IMAGE_TAG_STABLE:-}" \ | awk 'NF && !seen[$0]++' } @@ -743,6 +781,10 @@ find_matching_ci_base_ref() { remote_image_exists "${candidate}" || continue candidate_hash=$(get_remote_image_label "${candidate}" "vllm.ci_base.content_hash") if [[ "${candidate_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + if ! remote_ci_base_metadata_is_current "${candidate}"; then + echo "Found matching ci_base content hash but stale metadata: ${candidate}" >&2 + continue + fi printf '%s\n' "${candidate}" return 0 fi @@ -817,6 +859,10 @@ maybe_skip_existing_image() { if [[ -n "${remote_hash}" ]]; then echo "Remote ci_base content hash: ${remote_hash:0:16}..." if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + if ! remote_ci_base_metadata_is_current "${IMAGE_TAG}"; then + echo "Content hashes match but ci_base metadata is stale; rebuilding to refresh metadata" + return 0 + fi if ! refresh_ci_base_tags_from_ref "${IMAGE_TAG}"; then echo "ci_base tag refresh failed; rebuilding to push expected tags" return 0 @@ -998,12 +1044,104 @@ prepare_git_cache_metadata() { fi } +ci_base_metadata_pairs() { + local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" + local stages="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + local content_files="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" + local content_files_hash="" + local base_image="" + local base_image_digest="" + local git_branch="" + local -a content_paths=() + local -a content_args=() + + read -r -a content_paths <<< "${content_files}" + if [[ ${#content_paths[@]} -gt 0 ]]; then + content_files_hash=$(compute_content_hash "${content_paths[@]}") + fi + mapfile -t content_args < <( + get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}" + ) + + base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE") + if [[ -n "${base_image}" ]]; then + base_image_digest=$(resolve_image_digest "${base_image}") + fi + git_branch="${BUILDKITE_BRANCH:-${VLLM_BRANCH:-}}" + + metadata_pair "vllm.ci_base.metadata_version" "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" + metadata_pair "vllm.ci_base.content_hash" "${CI_BASE_CONTENT_HASH:-}" + metadata_pair "vllm.ci_base.content_files_hash" "${content_files_hash}" + metadata_pair "vllm.ci_base.content_files" "${content_files}" + metadata_pair "vllm.ci_base.content_args" "$(join_words "${content_args[@]}")" + metadata_pair "vllm.ci_base.dockerfile" "${dockerfile}" + metadata_pair "vllm.ci_base.dockerfile_stages" "${stages}" + metadata_pair "vllm.ci_base.image.primary" "${CI_BASE_IMAGE_TAG:-}" + metadata_pair "vllm.ci_base.image.content" "${CI_BASE_IMAGE_TAG_CONTENT_REF:-${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}}" + metadata_pair "vllm.ci_base.image.commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}}" + metadata_pair "vllm.ci_base.image.stable" "${CI_BASE_IMAGE_TAG_STABLE:-}" + metadata_pair "vllm.ci_base.git_commit" "${BUILDKITE_COMMIT:-}" + metadata_pair "vllm.ci_base.git_branch" "${git_branch}" + metadata_pair "vllm.ci_base.vllm_branch" "${VLLM_BRANCH:-}" + metadata_pair "vllm.ci_base.stable_branch" "${CI_BASE_STABLE_BRANCH:-main}" + + metadata_pair "vllm.rocm.base_image" "${base_image}" + metadata_pair "vllm.rocm.base_image_digest" "${base_image_digest}" + metadata_pair "vllm.rocm.pytorch_rocm_arch" "${PYTORCH_ROCM_ARCH:-}" + metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")" + metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")" + metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")" + metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")" + metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}" + metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")" + metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}" + metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")" + metadata_pair "vllm.rocm.rocshmem_commit" "${ROCSHMEM_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_BRANCH")}" + metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")" + metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}" + metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")" + metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")" + metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}" + metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}" + metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}" + + metadata_pair "vllm.buildkite.build_number" "${BUILDKITE_BUILD_NUMBER:-}" + metadata_pair "vllm.buildkite.build_id" "${BUILDKITE_BUILD_ID:-}" +} + +write_ci_base_metadata_annotations() { + local metadata="$1" + local key="" + local value="" + local annotation="" + + [[ -n "${metadata}" ]] || return 0 + while IFS=$'\t' read -r key value; do + [[ -n "${key}" && -n "${value}" ]] || continue + annotation="manifest:${key}=${value}" + printf ' "%s",\n' "$(hcl_escape_string "${annotation}")" + done <<< "${metadata}" +} + +write_ci_base_metadata_labels() { + local metadata="$1" + local key="" + local value="" + + [[ -n "${metadata}" ]] || return 0 + while IFS=$'\t' read -r key value; do + [[ -n "${key}" && -n "${value}" ]] || continue + printf ' "%s" = "%s"\n' \ + "$(hcl_escape_string "${key}")" \ + "$(hcl_escape_string "${value}")" + done <<< "${metadata}" +} + write_ci_base_label_override() { local target_name="" + local metadata="" local -a ci_base_targets=() - BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") - if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then return 0 fi @@ -1019,16 +1157,23 @@ write_ci_base_label_override() { return 0 fi + metadata=$(ci_base_metadata_pairs) + : > "${CI_BASE_LABEL_OVERRIDE_PATH}" for target_name in "${ci_base_targets[@]}"; do cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <> "${CI_BASE_LABEL_OVERRIDE_PATH}" + cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <> "${CI_BASE_LABEL_OVERRIDE_PATH}" + cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" < Date: Sun, 28 Jun 2026 14:09:18 +0800 Subject: [PATCH 083/138] [Model] Support Unlimited OCR (#46564) Signed-off-by: Tianyu Guo Signed-off-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Roger Wang --- docs/models/supported_models.md | 1 + tests/models/registry.py | 3 + .../core/test_single_type_kv_cache_manager.py | 52 +++- vllm/config/model.py | 4 + vllm/config/model_arch.py | 3 + .../layers/attention/__init__.py | 2 + .../layers/attention/rswa_attention.py | 37 +++ vllm/model_executor/models/config.py | 134 ++++++++++ vllm/model_executor/models/deepseek_ocr.py | 6 +- vllm/model_executor/models/deepseek_v2.py | 37 ++- vllm/model_executor/models/registry.py | 1 + vllm/model_executor/models/unlimited_ocr.py | 250 ++++++++++++++++++ vllm/tokenizers/registry.py | 6 +- .../chat_templates/registry.py | 1 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + .../configs/unlimited_ocr.py | 35 +++ .../model_arch_config_convertor.py | 7 + .../processors/deepseek_ocr.py | 4 +- .../processors/unlimited_ocr.py | 46 ++++ vllm/v1/attention/backend.py | 8 + vllm/v1/attention/backends/flash_attn.py | 117 +++++++- vllm/v1/attention/backends/flex_attention.py | 154 ++++++++++- vllm/v1/core/kv_cache_coordinator.py | 12 +- vllm/v1/core/kv_cache_manager.py | 14 +- vllm/v1/core/sched/scheduler.py | 1 + vllm/v1/core/single_type_kv_cache_manager.py | 122 +++++++-- vllm/v1/kv_cache_interface.py | 40 +++ vllm/v1/worker/gpu/attn_utils.py | 2 + vllm/v1/worker/gpu/input_batch.py | 3 + vllm/v1/worker/gpu/model_runner.py | 17 ++ vllm/v1/worker/gpu/model_states/default.py | 1 + .../gpu/model_states/encoder_decoder.py | 1 + .../worker/gpu/model_states/mamba_hybrid.py | 1 + vllm/v1/worker/gpu_model_runner.py | 8 + 35 files changed, 1084 insertions(+), 49 deletions(-) create mode 100644 vllm/model_executor/layers/attention/rswa_attention.py create mode 100644 vllm/model_executor/models/unlimited_ocr.py create mode 100644 vllm/transformers_utils/configs/unlimited_ocr.py create mode 100644 vllm/transformers_utils/processors/unlimited_ocr.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b0e0e3ce9c4..70b81aed9cd 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -629,6 +629,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | โœ…๏ธŽ | | `Tarsier2ForConditionalGeneration`^ | Tarsier2 | T + IE+ + VE+ | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | โœ…๏ธŽ | | `UltravoxModel` | Ultravox | T + AE+ | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | โœ…๏ธŽ | โœ…๏ธŽ | +| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I+ | `baidu/Unlimited-OCR`, etc. | โœ…๏ธŽ | โœ…๏ธŽ | Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it! diff --git a/tests/models/registry.py b/tests/models/registry.py index be271ea0777..463ce44851b 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -825,6 +825,9 @@ _MULTIMODAL_EXAMPLE_MODELS = { "DeepseekOCR2ForCausalLM": _HfExamplesInfo( "deepseek-ai/DeepSeek-OCR-2", ), + "UnlimitedOCRForCausalLM": _HfExamplesInfo( + "baidu/Unlimited-OCR", + ), "DotsOCRForCausalLM": _HfExamplesInfo( "rednote-hilab/dots.ocr", trust_remote_code=True ), diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 7e960c2a6a3..609c1428d19 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -14,9 +14,14 @@ from vllm.v1.core.kv_cache_utils import ( ) from vllm.v1.core.single_type_kv_cache_manager import ( ChunkedLocalAttentionManager, + RSWAManager, SlidingWindowManager, ) -from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + RSWASpec, + SlidingWindowSpec, +) pytestmark = pytest.mark.cpu_test @@ -327,6 +332,51 @@ def test_sliding_window_remove_skipped_blocks(): assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:]) +def test_rswa_remove_skipped_blocks_gap_range(): + block_size = 4 + rswa_spec = RSWASpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + rswa_window=8, + ) + block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=4) + manager = RSWAManager( + rswa_spec, + block_pool=block_pool, + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=block_size, + ) + + null_block_id = block_pool.null_block.block_id + original_block_ids = list(range(1000, 1010)) + block_table = [ + KVCacheBlock(id_) if id_ != null_block_id else block_pool.null_block + for id_ in original_block_ids + ] + manager.req_to_blocks["test"] = block_table + + prefix_len = 16 + + # Without num_prompt_tokens, R-SWA does not evict gap blocks. + manager.remove_skipped_blocks("test", 28) + assert [b.block_id for b in block_table] == original_block_ids + + # Gap = block 4 only (tokens [16, 20) fall in the gap). + manager.remove_skipped_blocks("test", 28, num_prompt_tokens=prefix_len) + expected = original_block_ids.copy() + expected[4] = null_block_id + assert [b.block_id for b in block_table] == expected + + # Window moves: blocks 5 and 6 also enter the gap; block 4 is already null. + manager.remove_skipped_blocks("test", 36, num_prompt_tokens=prefix_len) + expected[5] = null_block_id + expected[6] = null_block_id + assert [b.block_id for b in block_table] == expected + + def test_get_num_blocks_to_allocate(): block_size = 2 sliding_window_spec = SlidingWindowSpec( diff --git a/vllm/config/model.py b/vllm/config/model.py index fecb26aa7e0..ef0600af54d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1251,6 +1251,10 @@ class ModelConfig: def is_mm_prefix_lm(self) -> bool: return self.model_arch_config.is_mm_prefix_lm + @property + def rswa_window(self) -> int | None: + return self.model_arch_config.rswa_window + def get_head_size(self) -> int: return self.model_arch_config.head_size diff --git a/vllm/config/model_arch.py b/vllm/config/model_arch.py index 0b99df22b88..0b4744de489 100644 --- a/vllm/config/model_arch.py +++ b/vllm/config/model_arch.py @@ -56,5 +56,8 @@ class ModelArchitectureConfig: is_mm_prefix_lm: bool """Whether the model uses image bidirectional attention.""" + rswa_window: int | None + """Reference Sliding Window Attention window size (None disables R-SWA).""" + derived_max_model_len_and_key: tuple[float, str | None] """Derived maximum model length and key from the hf config.""" diff --git a/vllm/model_executor/layers/attention/__init__.py b/vllm/model_executor/layers/attention/__init__.py index ca3574164d5..c9e477fb114 100644 --- a/vllm/model_executor/layers/attention/__init__.py +++ b/vllm/model_executor/layers/attention/__init__.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderA from vllm.model_executor.layers.attention.prefill_prefix_lm_attention import ( PrefillPrefixLMAttention, ) +from vllm.model_executor.layers.attention.rswa_attention import RSWAAttention from vllm.model_executor.layers.attention.static_sink_attention import ( StaticSinkAttention, ) @@ -26,5 +27,6 @@ __all__ = [ "MLAAttention", "MMEncoderAttention", "PrefillPrefixLMAttention", + "RSWAAttention", "StaticSinkAttention", ] diff --git a/vllm/model_executor/layers/attention/rswa_attention.py b/vllm/model_executor/layers/attention/rswa_attention.py new file mode 100644 index 00000000000..c982722ff8e --- /dev/null +++ b/vllm/model_executor/layers/attention/rswa_attention.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.config.vllm import VllmConfig +from vllm.model_executor.layers.attention import Attention +from vllm.v1.kv_cache_interface import KVCacheSpec, RSWASpec, get_kv_quant_mode + + +class RSWAAttention(Attention): + """Attention layer that reports ``RSWASpec`` as its KV cache spec. + + Drop-in replacement for the standard ``Attention`` layer when the model is + configured with Reference Sliding Window Attention (R-SWA, + ``rswa_window > 0``). The actual masking logic lives in the attention + backend (FlexAttention or FA4 mask_mod); this layer only overrides + ``get_kv_cache_spec`` so the KV cache manager instantiates ``RSWAManager`` + (instead of ``FullAttentionManager``) and can therefore evict "gap" blocks + to keep per-request KV memory bounded at O(prefix + window). + """ + + def __init__(self, *args, rswa_window: int, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._rswa_window = rswa_window + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + spec = super().get_kv_cache_spec(vllm_config) + if spec is None: + return None + return RSWASpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size_v, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + rswa_window=self._rswa_window, + ) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index ac676149868..5c2278deb77 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -54,6 +54,139 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig): hf_config.is_causal = not hf_config.use_bidirectional_attention +class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + """Configure Unlimited-OCR attention backends for R-SWA and vision. + + Backend selection โ€” controlled by the standard ``--attention-config`` + CLI argument (priority order): + + 1. ``--attention-config '{"backend": "FLASH_ATTN"}'`` + โ†’ FA4 + rswa_mask_mod. Exact token-level R-SWA. + ``flash_attn_version`` is forced to 4 if not already set (R-SWA + mask_mod requires FA4; FA3 cannot express it). Raises if FA4 is + not available on this device. + + 2. ``--attention-config '{"backend": "FLEX_ATTENTION"}'`` + โ†’ FlexAttention R-SWA via Triton block mask. + + 3. ``--attention-config '{"backend": "auto"}'`` (or omitted) + โ†’ Auto-detect: FA4 if available (H20/H100 SM90), else FlexAttention. + + Regardless of backend, prefix caching is disabled for this model: R-SWA + decode-phase KV is not a pure causal function of the prefix (so decode + blocks are not reusable), and single-turn image-led OCR prompts rarely + hit the prefix cache. + + Example โ€” force FlexAttention even on a machine with FA4:: + + vllm serve baidu/Unlimited-OCR \\ + --attention-config '{"backend": "FLEX_ATTENTION"}' + """ + from vllm.v1.attention.backends.registry import AttentionBackendEnum + from vllm.vllm_flash_attn import is_fa_version_supported + + attn_config = vllm_config.attention_config + fa4_available = is_fa_version_supported(4) + + # โ”€โ”€ step 1: resolve backend โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # None means the user did not explicitly specify a backend; auto-select. + if attn_config.backend is None: + attn_config.backend = ( + AttentionBackendEnum.FLASH_ATTN + if fa4_available + else AttentionBackendEnum.FLEX_ATTENTION + ) + logger.info( + "Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).", + attn_config.backend.value, + fa4_available, + ) + + # โ”€โ”€ step 2: configure the chosen backend โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + if attn_config.backend == AttentionBackendEnum.FLASH_ATTN: + if not fa4_available: + raise RuntimeError( + "Unlimited-OCR: --attention-config backend=FLASH_ATTN " + "requires FA4 (rswa_mask_mod), but FA4 is not available on " + "this device/installation. Use backend=FLEX_ATTENTION or " + "upgrade vllm-flash-attn." + ) + # On SM90 (H20), the default FA version is FA3 regardless of FA4 + # availability (FA4 is only auto-upgraded when head_size > 256). + # The R-SWA mask_mod requires FA4, so force the version globally. + if attn_config.flash_attn_version is None: + attn_config.flash_attn_version = 4 + elif attn_config.flash_attn_version < 4: + logger.warning( + "Unlimited-OCR: flash_attn_version=%d cannot express the " + "R-SWA mask_mod; upgrading to 4.", + attn_config.flash_attn_version, + ) + attn_config.flash_attn_version = 4 + logger.info( + "Unlimited-OCR: FlashAttention FA%d + rswa_mask_mod โ€” exact R-SWA.", + attn_config.flash_attn_version, + ) + + elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION: + logger.info( + "Unlimited-OCR: FlexAttention โ€” R-SWA via Triton block mask%s.", + "" + if not fa4_available + else ( + " (FA4 available but not used; pass backend=FLASH_ATTN to upgrade)" + ), + ) + + else: + raise ValueError( + f"Unlimited-OCR: unsupported attention backend " + f"{attn_config.backend!r} for R-SWA. " + "Use FLASH_ATTN (FA4) or FLEX_ATTENTION." + ) + + # R-SWA windows the *generated* tokens, so a decode-token's KV is not a + # pure causal function of the prefix and cannot be safely reused across + # requests via prefix caching. Only the prompt/image prefix is cacheable, + # but OCR is single-turn with image-led prompts that rarely share a + # prefix, so prefix caching brings little benefit while complicating the + # KV cache manager. Disable it for this model. + cache_config = vllm_config.cache_config + if cache_config.enable_prefix_caching: + cache_config.enable_prefix_caching = False + logger.info( + "Unlimited-OCR: disabling prefix caching (R-SWA decode KV is not " + "cacheable, and single-turn image-led prompts rarely hit the " + "prefix cache)." + ) + + mm_config = getattr(vllm_config.model_config, "multimodal_config", None) + if mm_config is not None: + if mm_config.mm_encoder_attn_backend is None: + mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN + elif mm_config.mm_encoder_attn_backend == AttentionBackendEnum.FLASHINFER: + logger.warning( + "Unlimited-OCR: FlashInfer is not supported for the vision " + "encoder (the CLIP stage runs full attention without " + "cu_seqlens); falling back to FlashAttention." + ) + mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN + + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + text_config = model_config.hf_config.text_config + text_config.architectures = ["DeepseekV2ForCausalLM"] + if getattr(model_config.hf_config, "rswa_window", None) is None: + model_config.hf_config.rswa_window = 128 + # Propagate rswa_window to text_config so that DeepseekAttention (which + # receives text_config as its vllm_config.model_config.hf_config via + # init_vllm_registered_model) can read it and create RSWAAttention. + rswa_window = model_config.hf_config.rswa_window + text_config.rswa_window = rswa_window + + class Gemma4Config(VerifyAndUpdateConfig): @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: @@ -703,6 +836,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, } diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 0e061d6c6b5..b811afafb0e 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -220,8 +220,10 @@ class DeepseekOCRProcessingInfo(BaseProcessingInfo): patch_size = 16 downsample_ratio = 4 - if CROP_MODE: - if image_width <= 640 and image_height <= 640: + # Use the caller-supplied `cropping` flag so that callers that disable + # crop mode for multi-image requests get a consistent token count. + if cropping: + if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE: crop_ratio = [1, 1] else: # find the closest aspect ratio to the target diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 9b08ca9825e..09960050c06 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -45,7 +45,7 @@ from vllm.distributed import ( ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention import Attention, RSWAAttention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe import ( FusedMoE, @@ -174,15 +174,28 @@ class DeepseekAttention(nn.Module): max_position=max_position_embeddings, rope_parameters=config.rope_parameters, ) - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) + rswa_window = getattr(vllm_config.model_config.hf_config, "rswa_window", None) + if rswa_window is not None: + self.attn = RSWAAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + rswa_window=rswa_window, + ) + else: + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) def forward( self, @@ -588,12 +601,12 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase): compilation_config.static_forward_context[prefix] = self def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - return MLAAttentionSpec( # Only has one vector instead of K + V + return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, - ) + ) # Only has one vector instead of K + V def forward(self): ... diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1f9e3a24fe4..dfc034729d8 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -358,6 +358,7 @@ _MULTIMODAL_MODELS = { "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"), "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"), "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"), + "UnlimitedOCRForCausalLM": ("unlimited_ocr", "UnlimitedOCRForCausalLM"), "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"), "Eagle2_5_VLForConditionalGeneration": ( "eagle2_5_vl", diff --git a/vllm/model_executor/models/unlimited_ocr.py b/vllm/model_executor/models/unlimited_ocr.py new file mode 100644 index 00000000000..06dc02512f1 --- /dev/null +++ b/vllm/model_executor/models/unlimited_ocr.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Unlimited-OCR model compatible with HuggingFace weights. + +Unlimited-OCR (``baidu/Unlimited-OCR``) shares +the exact DeepSeek-OCR (gundam, ``base_size=1024`` / ``image_size=640`` / crop) +vision stack: a DeepEncoder (SAM-ViT-B + CLIP-L) followed by a linear MLP +projector, with the same image-token tiling layout. The only difference is the +language backbone, which is a DeepSeek-V2 *MoE* (64 routed + 2 shared experts, +``first_k_dense_replace=1``) that uses plain multi-head attention +(``use_mla=False``, ``qk_nope_head_dim == qk_rope_head_dim == 0``) instead of +the dense MLA decoder used by DeepSeek-OCR. + +vLLM's ``DeepseekV2DecoderLayer`` already dispatches to the plain-MHA +``DeepseekAttention`` whenever ``qk_nope_head_dim == qk_rope_head_dim == 0`` and +builds the MoE blocks straight from the config, so the whole DeepSeek-OCR +multimodal wrapper can be reused verbatim. Model-specific config (language +backbone architecture, FlexAttention for R-SWA, vision encoder backend, and +``rswa_window``) is applied in ``UnlimitedOCRForCausalLMConfig``. + +Attention backend: the reference applies Reference Sliding Window Attention +(R-SWA) -- the prompt/image tokens form a globally-visible prefix while the +*generated* tokens additionally attend only a fixed sliding window (128) of +recent tokens. We reproduce this (Level 1: full KV cache + custom mask) by +forcing the language model onto the FlexAttention backend and installing an +R-SWA ``mask_mod``. FlexAttention is the only backend able to express the +"global prefix + sliding window" mask; FlashAttention-3 / Triton only support a +uniform window (and additionally crash or compute incorrectly on this decoder's +10-head, +head_dim-128 shape), and FlashInfer's paged decode exposes no custom mask. The +window size is published via ``model_config.rswa_window``, which the model +runner reads to plumb per-request prefix lengths into the FlexAttention mask. + +The *vision encoder* (DeepEncoder's CLIP stage, head_dim 64) is unaffected and +does not use R-SWA: it runs a single full-attention prefill pass. FlashAttention, +Triton and torch SDPA all produce correct, equally fast results; only FlashInfer +is incompatible (its ViT path asserts on the varlen cu_seqlens metadata that this +CLIP encoder never builds). We default the encoder to FlashAttention and +transparently fall back to it if FlashInfer is requested. + +To suppress repetition on long documents, use ``NGramPerReqLogitsProcessor`` from +this module (same request-level processor as DeepSeek-OCR) with:: + + SamplingParams( + temperature=0.0, + max_tokens=8192, + extra_args={"ngram_size": 35, "window_size": 128}, + ) + +Image processing +---------------- +Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR), i.e. +``dynamic_preprocess`` runs with ``max_num=32``. + +Multi-image requests fall back to non-crop mode: crop ("gundam") mode is only +used for single-image input. DeepSeek-OCR does *not* have this restriction. + +Because that fallback makes the per-image processor output depend on *how many* +images are in the request, it breaks the assumption behind vLLM's per-item +multimodal processing cache (``MultiModalProcessorOnlyCache``). We handle this +the same way ``DeepseekVL2MultiModalProcessor`` does: only the single-image case +(which always crops) is cached, while multi-image requests bypass the cache and +are recomputed fresh -- see ``_cached_apply_hf_processor`` below. This keeps the +processing cache consistent (verified by ``test_processing_correctness``). +""" + +import math +from collections.abc import Mapping, Sequence + +from vllm.config import VllmConfig +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalKwargsItems +from vllm.multimodal.parse import ( + ImageEmbeddingItems, + ImageProcessorItems, + ImageSize, + MultiModalDataItems, +) +from vllm.multimodal.processing import PromptReplacement, PromptUpdate +from vllm.multimodal.processing.context import TimingContext +from vllm.multimodal.processing.inputs import ProcessorInputs +from vllm.multimodal.processing.processor import MultiModalProcessingInfo +from vllm.transformers_utils.processors.deepseek_ocr import ( + BASE_SIZE, + CROP_MODE, + IMAGE_SIZE, + count_tiles, +) + +from .deepseek_ocr import ( + DeepseekOCRDummyInputsBuilder, + DeepseekOCRForCausalLM, + DeepseekOCRMultiModalProcessor, + DeepseekOCRProcessingInfo, + NGramPerReqLogitsProcessor, +) + +__all__ = [ + "NGramPerReqLogitsProcessor", + "UnlimitedOCRForCausalLM", +] + +# Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR). +_UNLIMITED_OCR_MAX_CROPS = 32 + + +class UnlimitedOCRProcessingInfo(DeepseekOCRProcessingInfo): + """ProcessingInfo for Unlimited-OCR: same as DeepSeek-OCR but with + max_crops=32 instead of 6. The higher crop count allows tiling very large + document pages into up to 32 640ร—640 patches (dynamic_preprocess max_num=32). + """ + + def get_hf_config(self): + from vllm.transformers_utils.configs.unlimited_ocr import UnlimitedOCRConfig + + return self.ctx.get_hf_config(UnlimitedOCRConfig) + + def get_hf_processor(self, **kwargs: object): + from vllm.transformers_utils.processors.unlimited_ocr import ( + UnlimitedOCRProcessor, + ) + + v1_processor_config = dict( + image_size=IMAGE_SIZE, + base_size=BASE_SIZE, + crop_mode=CROP_MODE, + strategy="v1", + max_crops=_UNLIMITED_OCR_MAX_CROPS, + ) + return self.ctx.get_hf_processor( + UnlimitedOCRProcessor, + **{**v1_processor_config, **kwargs}, + ) + + def get_num_image_tokens( + self, *, image_width: int, image_height: int, cropping: bool = True + ) -> int: + patch_size = 16 + downsample_ratio = 4 + + # Honour the caller-supplied `cropping` flag: multi-image callers pass + # cropping=False to match UnlimitedOCRProcessor.tokenize_with_images. + if cropping: + if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE: + crop_ratio = [1, 1] + else: + crop_ratio = count_tiles( + image_width, + image_height, + max_num=_UNLIMITED_OCR_MAX_CROPS, + image_size=IMAGE_SIZE, + ) + num_width_tiles, num_height_tiles = crop_ratio + else: + num_width_tiles = num_height_tiles = 1 + + h = w = math.ceil((BASE_SIZE // patch_size) / downsample_ratio) + h2 = w2 = math.ceil((IMAGE_SIZE // patch_size) / downsample_ratio) + + global_views_tokens = h * (w + 1) + if num_width_tiles > 1 or num_height_tiles > 1: + local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1) + else: + local_views_tokens = 0 + + return global_views_tokens + local_views_tokens + 1 + + def get_image_size_with_most_features(self) -> ImageSize: + # With max_crops=32, the widest possible grid is 4ร—8 (aspect ratio 1:2). + # A 2560ร—5120 image (4ร—640 ร— 8ร—640) selects exactly 4ร—8=32 tiles and + # produces the maximum token count. + return ImageSize(width=640 * 4, height=640 * 8) + + +class UnlimitedOCRMultiModalProcessor(DeepseekOCRMultiModalProcessor): + """Multimodal processor for Unlimited-OCR. + + Disables crop mode for multi-image requests (to stay consistent with + ``UnlimitedOCRProcessor.tokenize_with_images``), and -- since that makes the + per-image output depend on the request's image count -- bypasses the + per-item processing cache for multi-image requests, exactly like + ``DeepseekVL2MultiModalProcessor``. + + DeepSeek-OCR does *not* apply either of these. + """ + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + + image_token_id = hf_processor.image_token_id + assert isinstance(image_token_id, int) + + def get_replacement_unlimited_ocr(item_idx: int): + images = mm_items.get_items( + "image", (ImageEmbeddingItems, ImageProcessorItems) + ) + + if isinstance(images, ImageEmbeddingItems): + num_image_tokens = images.get_feature_size(item_idx) + else: + size = images.get_image_size(item_idx) + + # Disable crop mode for multi-image input. + # UnlimitedOCRProcessor.tokenize_with_images applies the same + # fallback, so both paths must agree on the effective crop flag. + effective_cropping = CROP_MODE and len(images) == 1 + + num_image_tokens = self.info.get_num_image_tokens( + image_width=size.width, + image_height=size.height, + cropping=effective_cropping, + ) + return [image_token_id] * num_image_tokens + + return [ + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_replacement_unlimited_ocr, + ) + ] + + def _cached_apply_hf_processor( + self, + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> tuple[list[int], MultiModalProcessingInfo, bool]: + # The processor logic differs for single-image (crop) vs multi-image + # (no crop) requests. The processing cache assumes per-item output is + # invariant of how many images are passed per prompt, so we only cache + # the single-image case and recompute multi-image requests fresh. + if inputs.mm_data_items.get_count("image", strict=False) > 1: + return self._apply_hf_processor(inputs, timing_ctx) + + return super()._cached_apply_hf_processor(inputs, timing_ctx) + + +@MULTIMODAL_REGISTRY.register_processor( + UnlimitedOCRMultiModalProcessor, + info=UnlimitedOCRProcessingInfo, + dummy_inputs=DeepseekOCRDummyInputsBuilder, +) +class UnlimitedOCRForCausalLM(DeepseekOCRForCausalLM): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index eb7f8b0cf0d..f90e427aee0 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -31,7 +31,11 @@ logger = init_logger(__name__) # 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", "step3p7"} +_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = { + "step3_vl", + "step3p7", + "unlimited-ocr", +} _VLLM_TOKENIZERS = { "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), diff --git a/vllm/transformers_utils/chat_templates/registry.py b/vllm/transformers_utils/chat_templates/registry.py index a5f9bdac200..ed744742903 100644 --- a/vllm/transformers_utils/chat_templates/registry.py +++ b/vllm/transformers_utils/chat_templates/registry.py @@ -29,6 +29,7 @@ _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: dict[str, ChatTemplatePath] = { "colpali": CHAT_TEMPLATES_DIR / "template_basic.jinja", "deepseek_ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja", "deepseek_ocr2": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja", + "unlimited-ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja", "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja", "fuyu": CHAT_TEMPLATES_DIR / "template_fuyu.jinja", "minicpmv": _get_minicpmv_chat_template_fallback, diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index d6366407247..654d11df30d 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -124,6 +124,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( laguna="LagunaConfig", lfm2_moe="Lfm2MoeConfig", tarsier2="Tarsier2Config", + **{"unlimited-ocr": "UnlimitedOCRConfig"}, ) _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators"} diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 021eb2ea419..871cb524900 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -73,6 +73,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "RadioConfig": "vllm.transformers_utils.configs.radio", "SpeculatorsConfig": "vllm.transformers_utils.configs.speculators", "UltravoxConfig": "vllm.transformers_utils.configs.ultravox", + "UnlimitedOCRConfig": "vllm.transformers_utils.configs.unlimited_ocr", "Step3VLConfig": "vllm.transformers_utils.configs.step3_vl", "Step3VisionEncoderConfig": "vllm.transformers_utils.configs.step3_vl", "Step3TextConfig": "vllm.transformers_utils.configs.step3_vl", @@ -147,6 +148,7 @@ __all__ = [ "RadioConfig", "SpeculatorsConfig", "UltravoxConfig", + "UnlimitedOCRConfig", "Step3VLConfig", "Step3VisionEncoderConfig", "Step3TextConfig", diff --git a/vllm/transformers_utils/configs/unlimited_ocr.py b/vllm/transformers_utils/configs/unlimited_ocr.py new file mode 100644 index 00000000000..99e50a03c7d --- /dev/null +++ b/vllm/transformers_utils/configs/unlimited_ocr.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Unlimited-OCR (baidu/Unlimited-OCR) reuses +# the DeepSeek-OCR multimodal layout (DeepEncoder = SAM-ViT-B + CLIP-L, a linear +# MLP projector and a DeepSeek-V2 text backbone). The only architectural +# difference is the language model, which is a DeepSeek-V2 *MoE* with plain +# multi-head attention (``use_mla=False``) instead of the dense MLA backbone. +# We therefore reuse ``DeepseekVLV2Config`` for parsing the nested config. + +from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config + + +class UnlimitedOCRConfig(DeepseekVLV2Config): + model_type = "unlimited-ocr" + + # An explicit ``__init__`` is required: Transformers v5 processes each + # concrete config class' ``__init__`` signature to build nested sub-configs, + # and an empty subclass (only overriding ``model_type``) would skip + # ``DeepseekVLV2Config.__init__``, leaving ``text_config`` unset. + def __init__( + self, + tile_tag: str = "2D", + global_view_pos: str = "head", + candidate_resolutions: tuple[tuple[int, int]] = ((384, 384),), + rswa_window: int = 128, + **kwargs, + ): + super().__init__( + tile_tag=tile_tag, + global_view_pos=global_view_pos, + candidate_resolutions=candidate_resolutions, + **kwargs, + ) + self.rswa_window = rswa_window diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 483f1d8be81..e372834d68d 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -310,6 +310,12 @@ class ModelArchConfigConvertorBase: return False return self.hf_config.model_type in MM_PREFIX_LM_MODELS + def rswa_window(self) -> int | None: + value = getattr(self.hf_config, "rswa_window", None) + if value is None: + return None + return int(value) + def derive_max_model_len_and_key(self) -> tuple[float, str | None]: derived_max_model_len = float("inf") possible_keys = [ @@ -360,6 +366,7 @@ class ModelArchConfigConvertorBase: quantization_config=self.get_quantization_config(), is_deepseek_mla=self.is_deepseek_mla(), is_mm_prefix_lm=self.is_mm_prefix_lm(), + rswa_window=self.rswa_window(), derived_max_model_len_and_key=self.derive_max_model_len_and_key(), ) diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py index 68a2b1aaaa0..618070b506f 100644 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ b/vllm/transformers_utils/processors/deepseek_ocr.py @@ -161,10 +161,12 @@ class DeepseekOCRProcessor(ProcessorMixin): image_size: int = IMAGE_SIZE, base_size: int = BASE_SIZE, strategy: Literal["v1", "v2"] = "v1", + max_crops: int = MAX_CROPS, **kwargs, ): self.image_size = image_size self.base_size = base_size + self.max_crops = max_crops # image token calculation strategy for # Deepseek-OCR and Deepseek-OCR-2 @@ -332,7 +334,7 @@ class DeepseekOCRProcessor(ProcessorMixin): crop_ratio = [1, 1] elif cropping: images_crop_raw, crop_ratio = dynamic_preprocess( - image, image_size=self.image_size + image, image_size=self.image_size, max_num=self.max_crops ) else: crop_ratio = [1, 1] diff --git a/vllm/transformers_utils/processors/unlimited_ocr.py b/vllm/transformers_utils/processors/unlimited_ocr.py new file mode 100644 index 00000000000..927f19d0f93 --- /dev/null +++ b/vllm/transformers_utils/processors/unlimited_ocr.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Image processor for Unlimited-OCR (baidu/Unlimited-OCR).""" + +from PIL import Image + +from vllm.logger import init_logger +from vllm.transformers_utils.processors.deepseek_ocr import DeepseekOCRProcessor + +logger = init_logger(__name__) + + +class UnlimitedOCRProcessor(DeepseekOCRProcessor): + """DeepseekOCRProcessor variant for Unlimited-OCR. + + The only behavioural difference from the base processor is a multi-image + safeguard: when more than one image is present, crop ("gundam") mode is + disabled. + + Because the effective crop flag then depends on *how many* images are in the + request, the per-item processing output is no longer invariant of sibling + images. ``UnlimitedOCRMultiModalProcessor`` accounts for this by bypassing + the multimodal processing cache for multi-image requests (see its + ``_cached_apply_hf_processor``), so the two paths stay consistent. + + DeepSeek-OCR does *not* have this restriction because its ``max_crops=6`` is + small enough to be safe for multi-image use. + """ + + def tokenize_with_images( + self, + conversation: str, + images: list[Image.Image], + bos: bool = True, + eos: bool = True, + cropping: bool = True, + ): + if len(images) > 1 and cropping: + logger.warning_once( + "Unlimited-OCR: crop mode is not supported for multi-image " + "input. Falling back to cropping=False." + ) + cropping = False + return super().tokenize_with_images( + conversation, images, bos=bos, eos=eos, cropping=cropping + ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index ccd70c6ca3c..61a4e521c40 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -455,6 +455,13 @@ class CommonAttentionMetadata: where bidirectional attention should apply. None for text-only batches or non-PrefixLM models.""" + rswa_prefix_lens: torch.Tensor | None = None + """(batch_size,) per-request prefix length (prompt/image token count) for + Reference Sliding Window Attention (R-SWA). Tokens with logical index below + this stay globally visible; later (generated) tokens additionally see a + fixed sliding window. None disables R-SWA. The attention backend copies this + into its own persistent buffer and reads ``rswa_window`` from model config.""" + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) _seq_lens_cpu: torch.Tensor | None = None _num_computed_tokens_cpu: torch.Tensor | None = None @@ -539,6 +546,7 @@ class CommonAttentionMetadata: dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), is_prefilling=maybe_slice_reqs(self.is_prefilling), + rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), ) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 75231bafeed..c167708ac9c 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -256,6 +256,16 @@ class FlashAttentionMetadata: # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range. mm_prefix_range_tensor: torch.Tensor | None = None + # Reference Sliding Window Attention (R-SWA) fields. + # rswa_prefix_lens: per-request prompt lengths [num_reqs], int32, CUDA. + # rswa_window: sliding window size (scalar int, for logic checks). + # rswa_window_tensor: [1] int32 CUDA tensor โ€” pre-allocated in build() so + # no CPUโ†’CUDA copy is needed inside forward() during CUDA graph capture. + # Only populated when the model uses R-SWA (Unlimited-OCR). + rswa_prefix_lens: torch.Tensor | None = None + rswa_window: int | None = None + rswa_window_tensor: torch.Tensor | None = None + def _get_sliding_window_configs( vllm_config: VllmConfig, @@ -386,6 +396,19 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad # populated on first build() call. self.aot_sliding_window: tuple[int, int] | None = None + # R-SWA: persistent CUDA-graph-safe buffers owned by this builder. + self.rswa_window: int | None = self.model_config.rswa_window + self.persistent_rswa_prefix_lens: torch.Tensor | None = None + self.persistent_rswa_window_tensor: torch.Tensor | None = None + if self.rswa_window is not None: + max_num_reqs = vllm_config.scheduler_config.max_num_seqs + self.persistent_rswa_prefix_lens = torch.zeros( + max_num_reqs, dtype=torch.int32, device=self.device + ) + self.persistent_rswa_window_tensor = torch.tensor( + [self.rswa_window], dtype=torch.int32, device=self.device + ) + def build( self, common_prefix_len: int, @@ -589,6 +612,22 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad mm_ranges, num_reqs, seq_lens.device ) + # R-SWA: copy prefix lengths into persistent buffers (outside the + # compiled region) so forward() never allocates during CUDA graph + # capture. rswa_window is a static model config scalar read here. + if ( + self.rswa_window is not None + and common_attn_metadata.rswa_prefix_lens is not None + ): + assert self.persistent_rswa_prefix_lens is not None + assert self.persistent_rswa_window_tensor is not None + src = common_attn_metadata.rswa_prefix_lens + rswa_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs] + rswa_prefix_lens.copy_(src[:num_reqs], non_blocking=True) + attn_metadata.rswa_prefix_lens = rswa_prefix_lens + attn_metadata.rswa_window = self.rswa_window + attn_metadata.rswa_window_tensor = self.persistent_rswa_window_tensor + return attn_metadata def update_block_table( @@ -805,7 +844,7 @@ class FlashAttentionImpl(AttentionImpl): ) return output else: - sliding_window_size = ( + sliding_window_size: list[int] | None = ( list(self.sliding_window) if self.sliding_window is not None else None @@ -840,6 +879,25 @@ class FlashAttentionImpl(AttentionImpl): mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges) mm_aux = [mm_prefix_ranges] + # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level + # mask without block-size approximation. The mask_mod encodes + # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which + # supersedes any FA-layer sliding_window_size parameter. + rswa_mask_mod_fn = None + rswa_aux = None + if ( + attn_metadata.rswa_prefix_lens is not None + and self.vllm_flash_attn_version == 4 + and not is_dynamic_causal + ): + rswa_mask_mod_fn = _make_rswa_mask_mod() + rswa_aux = [ + attn_metadata.rswa_prefix_lens.to(torch.int32), + attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor + ] + # mask_mod fully expresses R-SWA; disable FA's own window. + sliding_window_size = None + dynamic_causal = None if isinstance(causal, torch.Tensor): if self.vllm_flash_attn_version != 4: @@ -873,8 +931,8 @@ class FlashAttentionImpl(AttentionImpl): dynamic_causal=dynamic_causal, num_splits=attn_metadata.max_num_splits, s_aux=self.sinks, - mask_mod=mm_mask_mod, - aux_tensors=mm_aux, + mask_mod=rswa_mask_mod_fn or mm_mask_mod, + aux_tensors=rswa_aux or mm_aux, ) return output @@ -1152,6 +1210,59 @@ def _make_mm_prefix_mask_mod(max_ranges: int): return mm_prefix_mask_mod +def _make_rswa_mask_mod(): + """Build a CuTE-DSL mask_mod for Reference Sliding Window Attention (R-SWA). + + FA4 varlen + paged-KV convention (verified from cute/mask.py apply_mask): + q_idx = LOCAL query-token offset (0 .. seqlen_q - 1) within this sequence. + kv_idx = LOCAL KV-token position (0 .. seqlen_k - 1) within this sequence. + + To recover the ABSOLUTE token position (needed for causal and the sliding + window distance), use the standard offset: + abs_q = q_idx + (seqlen_k - seqlen_q) + + R-SWA keep condition: + abs_q >= kv_idx (causal: KV at or before the query) + AND (kv_idx < prefix_len (global prefix is always visible) + OR abs_q - kv_idx < window) (generated tokens: sliding window) + + aux_tensors[0]: prefix_lens [num_reqs] int32 โ€” per-request prefill length. + aux_tensors[1]: rswa_window [1] int32 โ€” decode sliding window size. + + use_fast_sampling=True lets FA4 skip fully-masked KV blocks (gap blocks) + without loading their data. + """ + import cutlass.cute as cute + from cutlass import Int32 # type: ignore[attr-defined] + + from vllm.vllm_flash_attn.cute.utils import ( # type: ignore[import-untyped] + scalar_to_ssa, + ) + + @cute.jit + def rswa_mask_mod( + batch_idx: cute.TensorSSA, + head_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info, + aux_tensors, + ): + b = batch_idx[0] + prefix_len = scalar_to_ssa(aux_tensors[0][b], Int32) + window = scalar_to_ssa(aux_tensors[1][0], Int32) + # Convert local q offset to absolute token position. + offset = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32) + abs_q = q_idx + offset + causal = kv_idx <= abs_q + in_prefix = kv_idx < prefix_len + in_window = (abs_q - kv_idx) < window + return causal & (in_prefix | in_window) + + rswa_mask_mod.use_fast_sampling = True + return rswa_mask_mod + + def use_cascade_attention( common_prefix_len: int, query_lens: np.ndarray, diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 983544b5602..c45294bfc79 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -408,6 +408,11 @@ class FlexAttentionMetadata: sliding_window: int | None = None mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None block_sparsity_hint: BlockSparsityHint | None = None + # Reference Sliding Window Attention (R-SWA): per-request prefix length + # (prompt/image tokens stay globally visible) plus a sliding window over + # generated tokens. Both must be set to enable. + rswa_prefix_lens: torch.Tensor | None = None + rswa_window: int | None = None @cached_property def logical_block_ids(self): @@ -571,6 +576,52 @@ class FlexAttentionMetadata: return final_mask_mod + def get_rswa_mask_mod(self) -> _mask_mod_signature: + """Creates the Reference Sliding Window Attention (R-SWA) mask_mod. + + R-SWA keeps the whole prefix (image + prompt tokens, i.e. logical index + ``< prefix_len``) globally visible while generated tokens additionally + attend a fixed sliding window of recent tokens. This term is combined + with the base causal mask via logical AND, so it only ever *removes* + far-away generated tokens that fall outside the window and outside the + prefix. + """ + + assert self.doc_ids is not None + assert self.rswa_prefix_lens is not None + assert self.rswa_window is not None + doc_ids = self.doc_ids + prefix_lens = self.rswa_prefix_lens + window = self.rswa_window + + def rswa_mask_mod( + q_req: torch.Tensor, + logical_q_idx: torch.Tensor, + logical_kv_idx: torch.Tensor, + ) -> torch.Tensor: + prefix_len = prefix_lens[q_req] + in_prefix = logical_kv_idx < prefix_len + in_window = (logical_q_idx - logical_kv_idx) < window + return in_prefix | in_window + + def final_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + physical_kv_idx: torch.Tensor, + ) -> torch.Tensor: + (is_valid, logical_q_idx, logical_kv_idx) = ( + self._convert_physical_to_logical(doc_ids, q_idx, physical_kv_idx) + ) + q_req = doc_ids[q_idx] + return torch.where( + is_valid, + rswa_mask_mod(q_req, logical_q_idx, logical_kv_idx), + False, + ) + + return final_mask_mod + def get_mask_mod(self): # Stage-1: initialize the base mask_mod # (causal mask for decoder or bidirectional mask for encoder) @@ -588,6 +639,10 @@ class FlexAttentionMetadata: # Add prefix LM mask for vision-language prefix LM attention prefix_lm_mask_mod = self.get_prefix_lm_mask_mod() mask_mod = or_masks(mask_mod, prefix_lm_mask_mod) + if self.rswa_window is not None and self.rswa_prefix_lens is not None: + # Reference Sliding Window Attention: AND with the base causal mask + # (prefix stays global, generated tokens use a sliding window). + mask_mod = and_masks(mask_mod, self.get_rswa_mask_mod()) return mask_mod def get_transformed_score_mod(self) -> _score_mod_signature | None: @@ -663,9 +718,21 @@ class FlexAttentionMetadata: self.doc_ids, : cdiv(self.max_seq_len, self.block_size) ] - custom_hint = self.block_sparsity_hint is not None + # block_table slots beyond each request's seq_len may contain garbage + # physical page ids (see physical_to_logical_mapping). With batched + # decode, max_seq_len is the batch max while shorter requests still + # index all columns up to that max unless masked here. + num_blocks = self.num_blocks_per_seq[self.doc_ids] + past_seq = self.logical_block_ids[None, :] >= num_blocks[:, None] + used_pages.masked_fill_(past_seq, 0) - if self.sliding_window or custom_hint: + custom_hint = self.block_sparsity_hint is not None + use_rswa = self.rswa_window is not None and self.rswa_prefix_lens is not None + needs_per_q_pruning = ( + self.causal or self.sliding_window or custom_hint or use_rswa + ) + + if needs_per_q_pruning: device = used_pages.device assert self.doc_ids is not None token_indices = torch.arange( @@ -676,6 +743,12 @@ class FlexAttentionMetadata: - self.query_start_loc[self.doc_ids] + self.decode_offset[self.doc_ids] ) + block_starts = self.logical_block_ids * self.block_size + block_ends = block_starts + self.block_size + + if self.causal: + future_blocks = block_starts[None, :] > logical_q_idx[:, None] + used_pages.masked_fill_(future_blocks, 0) if self.sliding_window: assert self.sliding_window is not None @@ -685,6 +758,23 @@ class FlexAttentionMetadata: min_block_idx = min_kv_idx // self.block_size sliding_mask = self.logical_block_ids >= min_block_idx[:, None] used_pages.masked_fill_(~sliding_mask, 0) + if use_rswa: + # R-SWA keeps prefix KV globally visible and applies a sliding + # window over generated tokens. Prune blocks that fall entirely + # in the "hole" between prefix_len and the current window so + # FlexAttention does not gather invalid paged-KV slots (this + # mirrors uniform sliding-window block pruning above). + assert self.rswa_prefix_lens is not None + assert self.rswa_window is not None + prefix_len = self.rswa_prefix_lens[self.doc_ids] + min_kv_window = torch.maximum( + prefix_len, + logical_q_idx - (self.rswa_window - 1), + ) + in_gap = (block_starts[None, :] >= prefix_len[:, None]) & ( + block_ends[None, :] <= min_kv_window[:, None] + ) + used_pages.masked_fill_(in_gap, 0) if custom_hint: assert self.block_sparsity_hint is not None q_block_idx = logical_q_idx // self.block_size @@ -798,12 +888,36 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat self.max_num_query_groups = cdiv(max_num_batched_tokens, self.q_block_size) max_num_pages_per_seq = cdiv(self.max_model_len, self.block_size) self.max_num_kv_indices = self.q_block_size * max_num_pages_per_seq + # R-SWA uses q_block_size=1 so block lists are not merged across requests + # in a q-group (mixed-length batches otherwise gather foreign paged-KV). + self.max_num_rswa_query_groups = max_num_batched_tokens + # +1 sentinel column: the flex-attention kernel's get_offset_for_next_block + # always prefetches kv_indices[q, kv_num_blocks] (one past the last valid + # entry) to compute the jump offset for the next loop iteration. When + # kv_num_blocks[q] == W (every page of the sequence is live), that prefetch + # reads column W of the persistent buffer. Without the extra column this + # would land on stale data from a previous step (the buffer is wider than W + # but is never fully zeroed), producing an out-of-bounds K/V pointer and a + # CUDA illegal memory access. Allocating W_max+1 columns and initialising + # the whole buffer to -1 ensures the sentinel slot is always safe to read. + self.max_num_rswa_kv_indices = max_num_pages_per_seq + 1 self.persistent_kv_num_blocks = torch.empty( self.max_num_query_groups, dtype=torch.int32, device=device ) + self.persistent_rswa_kv_num_blocks = torch.empty( + self.max_num_rswa_query_groups, dtype=torch.int32, device=device + ) self.persistent_offset_tensor = torch.empty( max_num_seqs, dtype=torch.int32, device=device ) + # Persistent buffer for R-SWA per-request prefix lengths so the device + # address stays stable across steps (required for CUDA graph replay). + self.rswa_window: int | None = self.model_config.rswa_window + self.persistent_rswa_prefix_lens: torch.Tensor | None = None + if self.rswa_window is not None: + self.persistent_rswa_prefix_lens = torch.empty( + max_num_seqs, dtype=torch.int32, device=device + ) self.persistent_doc_ids = torch.empty( max_num_batched_tokens, dtype=torch.int32, device=device ) @@ -811,6 +925,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat # initialize later when we can access block_table self.persistent_physical_to_logical = None self.persistent_kv_indices = None + self.persistent_rswa_kv_indices = None self.custom_logical_mask_mod: _mask_mod_signature | None = None if self._uses_full_cudagraphs(): @@ -936,6 +1051,26 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat dtype=torch.int32, device=self.device, ) + if self.persistent_rswa_kv_indices is None: + # Initialise to -1 so the +1 sentinel column (see max_num_rswa_kv_indices) + # is always a safe pad value for the flex kernel's prefetch. + self.persistent_rswa_kv_indices = torch.full( + (self.max_num_rswa_query_groups, self.max_num_rswa_kv_indices), + fill_value=-1, + dtype=torch.int32, + device=self.device, + ) + + use_rswa = self.rswa_window is not None + q_block_size = 1 if use_rswa else self.q_block_size + persistent_kv_indices = ( + self.persistent_rswa_kv_indices if use_rswa else self.persistent_kv_indices + ) + persistent_kv_num_blocks = ( + self.persistent_rswa_kv_num_blocks + if use_rswa + else self.persistent_kv_num_blocks + ) inverse_block_table = copy_to_persistent( self.persistent_physical_to_logical, inverse_block_table @@ -944,6 +1079,13 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat offset_tensor = common_attn_metadata.compute_num_computed_tokens() offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor) + rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens + if use_rswa and rswa_prefix_lens is not None: + assert self.persistent_rswa_prefix_lens is not None + rswa_prefix_lens = copy_to_persistent( + self.persistent_rswa_prefix_lens, rswa_prefix_lens + ) + uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec) logical_mask_mod = ( bidirectional_mask_mod @@ -986,12 +1128,14 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat # attention block mask for encoder-only models, disable it temporarily. # see: https://github.com/vllm-project/vllm/pull/27329#issuecomment-3431484053 direct_build=self.direct_build and uses_paged_kv, - q_block_size=self.q_block_size, + q_block_size=q_block_size, kv_block_size=self.kv_block_size, - persistent_kv_indices=self.persistent_kv_indices, - persistent_kv_num_blocks=self.persistent_kv_num_blocks, + persistent_kv_indices=persistent_kv_indices, + persistent_kv_num_blocks=persistent_kv_num_blocks, persistent_doc_ids=self.persistent_doc_ids, mm_prefix_range=common_attn_metadata.mm_req_doc_ranges, + rswa_prefix_lens=rswa_prefix_lens, + rswa_window=self.rswa_window, ) # Pre-build block_mask so it is ready before CUDA graph capture. diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 48f597e1f24..a759d7a80ad 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -329,7 +329,10 @@ class KVCacheCoordinator(ABC): ] def remove_skipped_blocks( - self, request_id: str, total_computed_tokens: int + self, + request_id: str, + total_computed_tokens: int, + num_prompt_tokens: int | None = None, ) -> None: """ Remove the blocks that are no longer needed from `blocks` and replace @@ -339,9 +342,14 @@ class KVCacheCoordinator(ABC): request_id: The request ID. total_computed_tokens: The total number of computed tokens, including local computed tokens and external computed tokens. + num_prompt_tokens: Optional prompt length. R-SWA managers use this to + free gap blocks between the prefill tail and decode window; other + manager types ignore it. """ for manager in self.single_type_managers: - manager.remove_skipped_blocks(request_id, total_computed_tokens) + manager.remove_skipped_blocks( + request_id, total_computed_tokens, num_prompt_tokens + ) def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]: """ diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index b0f6655bf95..57cd1490e81 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -398,7 +398,9 @@ class KVCacheManager: # Should call this function before allocating new blocks to reduce # the number of evicted blocks. self.coordinator.remove_skipped_blocks( - request.request_id, total_computed_tokens + request.request_id, + total_computed_tokens, + num_prompt_tokens=request.num_prompt_tokens, ) num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( @@ -468,7 +470,10 @@ class KVCacheManager: self.coordinator.free(request.request_id) def remove_skipped_blocks( - self, request_id: str, total_computed_tokens: int + self, + request_id: str, + total_computed_tokens: int, + num_prompt_tokens: int | None = None, ) -> None: """Remove the blocks that are no longer needed from `blocks` and replace the removed blocks with null_block. @@ -477,8 +482,11 @@ class KVCacheManager: request_id: The request ID. total_computed_tokens: The total number of computed tokens, including local computed tokens and external computed tokens. + num_prompt_tokens: Optional prompt length for R-SWA gap eviction. """ - self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens) + self.coordinator.remove_skipped_blocks( + request_id, total_computed_tokens, num_prompt_tokens + ) def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: """Pop the request's bookkeeping and return its blocks without diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ab9fd5e3433..ec479f09304 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2342,6 +2342,7 @@ class Scheduler(SchedulerInterface): self.kv_cache_manager.remove_skipped_blocks( request_id=request.request_id, total_computed_tokens=request.num_computed_tokens, + num_prompt_tokens=request.num_prompt_tokens, ) block_ids = self.kv_cache_manager.get_block_ids(request.request_id) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index e21c20a2281..642fe3e6a08 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -20,6 +20,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, MambaSpec, MLAAttentionSpec, + RSWASpec, SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, @@ -476,8 +477,38 @@ class SingleTypeKVCacheManager(ABC): raise NotImplementedError + def _remove_blocks_in_range( + self, + request_id: str, + first_block: int, + last_block: int, + ) -> None: + """Free blocks in ``[first_block, last_block)`` and replace with null_block. + + Iterates backward so newly-evictable tail blocks are reached even after + earlier blocks in the range were nulled in a prior call. + """ + if request_id not in self.req_to_blocks: + return + if first_block >= last_block: + return + blocks = self.req_to_blocks[request_id] + last_block = min(last_block, len(blocks)) + + freed: list[KVCacheBlock] = [] + for i in range(last_block - 1, first_block - 1, -1): + if blocks[i] == self._null_block: + break + freed.append(blocks[i]) + blocks[i] = self._null_block + if freed: + self.block_pool.free_blocks(freed) + def remove_skipped_blocks( - self, request_id: str, total_computed_tokens: int + self, + request_id: str, + total_computed_tokens: int, + num_prompt_tokens: int | None = None, ) -> None: """ Remove and free the blocks that are no longer needed for attention computation. @@ -490,7 +521,11 @@ class SingleTypeKVCacheManager(ABC): request_id: The request ID. total_computed_tokens: The total number of computed tokens, including local computed tokens and external computed tokens. + num_prompt_tokens: Optional prompt length for attention types (e.g. + R-SWA) that evict a middle gap rather than a head prefix. Ignored + by the default implementation. """ + del num_prompt_tokens # Remove the blocks that will be skipped during attention computation. num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) if num_skipped_tokens <= 0: @@ -506,18 +541,7 @@ class SingleTypeKVCacheManager(ABC): # range), so we must cap to the number of blocks that currently exist for # this request. num_skipped_blocks = min(num_skipped_blocks, len(blocks)) - removed_blocks: list[KVCacheBlock] = [] - # Because the block starts from index 0, the num_skipped_block-th block - # corresponds to index num_skipped_blocks - 1. - for i in range(num_skipped_blocks - 1, -1, -1): - if blocks[i] == self._null_block: - # If the block is already a null block, the blocks before it - # should also have been set to null blocks by the previous calls - # to this function. - break - removed_blocks.append(blocks[i]) - blocks[i] = self._null_block - self.block_pool.free_blocks(removed_blocks) + self._remove_blocks_in_range(request_id, 0, num_skipped_blocks) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: """ @@ -598,6 +622,50 @@ class FullAttentionManager(SingleTypeKVCacheManager): return num_common_blocks +class RSWAManager(FullAttentionManager): + """KV cache manager for Reference Sliding Window Attention (R-SWA). + + When ``num_prompt_tokens`` is supplied to ``remove_skipped_blocks``, frees + gap blocks between the prefill tail and the current decode window. This + bounds per-request KV memory at O(prefix_len + rswa_window) instead of + growing linearly with decode length. + """ + + def __init__(self, kv_cache_spec: RSWASpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.rswa_window: int = kv_cache_spec.rswa_window + + def remove_skipped_blocks( + self, + request_id: str, + total_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + """Free gap blocks that are no longer needed for attention. + + Gap = blocks entirely within + [ceil(prefix_len / block_size) * block_size, + max(prefix_len, total_computed_tokens - rswa_window)) + + Freed blocks are replaced with null_block in req_to_blocks so the + block_table passed to FA4 is valid (null_block KV is all-zero; + rswa_mask_mod marks gap positions as non-visible so FA4 skips them). + """ + if num_prompt_tokens is None: + super().remove_skipped_blocks( + request_id, total_computed_tokens, num_prompt_tokens + ) + return + + bs = self.block_size + # First block fully after the prefill boundary. + first_gap_block = cdiv(num_prompt_tokens, bs) + # Decode window start position; blocks before this are evictable. + window_start = max(num_prompt_tokens, total_computed_tokens - self.rswa_window) + last_gap_block = window_start // bs # exclusive upper bound + self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block) + + class SlidingWindowManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) @@ -1072,7 +1140,12 @@ class MambaManager(SingleTypeKVCacheManager): return mask - def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None: + def remove_skipped_blocks( + self, + request_id: str, + num_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: assert isinstance(self.kv_cache_spec, MambaSpec) # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain @@ -1082,7 +1155,9 @@ class MambaManager(SingleTypeKVCacheManager): # that we might actually need. num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks) - super().remove_skipped_blocks(request_id, num_computed_tokens) + super().remove_skipped_blocks( + request_id, num_computed_tokens, num_prompt_tokens + ) if self.mamba_cache_mode == "align": # `last_state_block_idx` refers to the block index allocated two steps ago. # The block allocated in the previous step is used to copy Mamba states @@ -1401,10 +1476,16 @@ def get_manager_for_kv_cache_spec( assert manager_class is not None, ( f"No manager registered for KVCacheSpec {type(kv_cache_spec)}" ) - # SlidingWindow / ChunkedLocalAttention managers recycle blocks across - # chunks; the runtime admission cap must match the recycling-aware bound - # the startup pool sizer uses (single source of truth: the spec method). - if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): + # SlidingWindow / ChunkedLocalAttention managers recycle blocks; + # the runtime admission cap must match the recycling-aware bound the + # startup pool sizer uses (single source of truth: the spec method). + # R-SWA also recycles gap blocks but peak physical KV still fits the + # full-attention bound (prefix + window <= max_model_len), so it inherits + # FullAttentionSpec sizing without a separate admission cap. + if isinstance( + kv_cache_spec, + (SlidingWindowSpec, ChunkedLocalAttentionSpec), + ): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( max_num_batched_tokens=max_num_batched_tokens, @@ -1457,6 +1538,9 @@ def register_all_kvcache_specs(vllm_config): KVCacheSpecRegistry.register( MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec ) + KVCacheSpecRegistry.register( + RSWASpec, RSWAManager, uniform_type_base_spec=FullAttentionSpec + ) # NOTE(Mengqing): HiddenStateCacheSpec won't take part in # grouping, thus the uniform_type_base_spec is just a # placeholder. diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index b312a0fbeef..323b1e763a5 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -437,6 +437,46 @@ class HiddenStateCacheSpec(MLAAttentionSpec): pass +@dataclass(frozen=True, kw_only=True) +class RSWASpec(FullAttentionSpec): + """KV cache spec for Reference Sliding Window Attention (R-SWA). + + Prefill (image + text prompt) tokens are always globally visible. + Only the last ``rswa_window`` generated tokens are kept in the KV cache; + gap blocks (between the prefill tail and the current decode window) are + evicted during each decode step to bound memory at + O(prefix_blocks + window_blocks). + """ + + rswa_window: int + + @classmethod + def merge(cls, specs: list[RSWASpec]) -> RSWASpec: + assert all(isinstance(spec, RSWASpec) for spec in specs), ( + "All attention layers in the same KV cache group must be RSWASpec." + ) + rswa_windows = {spec.rswa_window for spec in specs} + assert len(rswa_windows) == 1, ( + f"All R-SWA layers must share the same rswa_window, got {rswa_windows}" + ) + # Delegate common field merging to the parent, then reattach rswa_window. + base = FullAttentionSpec.merge(specs) # type: ignore[arg-type] + return cls( + block_size=base.block_size, + num_kv_heads=base.num_kv_heads, + head_size=base.head_size, + head_size_v=base.head_size_v, + dtype=base.dtype, + kv_quant_mode=base.kv_quant_mode, + page_size_padded=base.page_size_padded, + indexes_kv_by_block_stride=base.indexes_kv_by_block_stride, + sliding_window=base.sliding_window, + attention_chunk_size=base.attention_chunk_size, + non_causal=base.non_causal, + rswa_window=rswa_windows.pop(), + ) + + @dataclass(frozen=True, kw_only=True) class ChunkedLocalAttentionSpec(AttentionSpec): attention_chunk_size: int diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 737feb7d277..758bd3bac7a 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -469,6 +469,7 @@ def build_attn_metadata( model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, for_cudagraph_capture: bool = False, causal: bool = True, + rswa_prefix_lens: torch.Tensor | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -501,6 +502,7 @@ def build_attn_metadata( causal=causal, dcp_local_seq_lens=dcp_local_seq_lens, positions=positions, + rswa_prefix_lens=rswa_prefix_lens, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index d745dc6abf9..a6a2b296e38 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -96,6 +96,9 @@ class InputBatch: # Whether any requests in batch use structured output. has_structured_output_reqs: bool + # [num_reqs_after_padding] per-request prompt length for R-SWA (optional). + rswa_prefix_lens: torch.Tensor | None = None + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 927ece4fbac..ce1bb7f5504 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -223,6 +223,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_num_tokens=self.max_num_tokens, device=self.device, ) + # R-SWA: persistent GPU buffer for per-request prefix lengths (CUDA-graph safe). + self.rswa_prefix_lens_buffer: torch.Tensor | None = None + if self.model_config.rswa_window is not None: + self.rswa_prefix_lens_buffer = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) if self.use_pp: self.pp_handler = PPHandler( @@ -985,6 +991,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.use_pp: # max_seq_len is only consumed by the PP `compute_need_sampled_mask` max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np] + + rswa_prefix_lens = None + if self.rswa_prefix_lens_buffer is not None: + rswa_prefix_lens = self.rswa_prefix_lens_buffer[:num_reqs_padded] + rswa_prefix_lens[:num_reqs] = self.req_states.prompt_len.gpu[ + idx_mapping[:num_reqs] + ] + if num_reqs_padded > num_reqs: + rswa_prefix_lens[num_reqs:].zero_() + return InputBatch( req_ids=req_ids, num_reqs=num_reqs, @@ -1015,6 +1031,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=scheduler_output.has_structured_output_requests, + rswa_prefix_lens=rswa_prefix_lens, ) def prepare_attn( diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 2e14eb2e7d9..f760fc36dea 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -168,5 +168,6 @@ class DefaultModelState(ModelState): dcp_local_seq_lens=input_batch.dcp_local_seq_lens, positions=input_batch.positions, for_cudagraph_capture=for_capture, + rswa_prefix_lens=input_batch.rswa_prefix_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index 889e624623d..9edda27538e 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -146,6 +146,7 @@ class EncoderDecoderModelState(ModelState): dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=enc_dec_attn_metadata, for_cudagraph_capture=for_capture, + rswa_prefix_lens=input_batch.rswa_prefix_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 329f008a4e3..e08b09f1895 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -141,6 +141,7 @@ class MambaHybridModelState(DefaultModelState): dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=mamba_attn_metadata, for_cudagraph_capture=for_capture, + rswa_prefix_lens=input_batch.rswa_prefix_lens, ) def postprocess_state( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a823d9..6af53115775 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2335,6 +2335,13 @@ class GPUModelRunner( req_idx = self.input_batch.req_id_to_index[req_id] req_doc_ranges[req_idx] = image_doc_ranges + # Reference Sliding Window Attention (R-SWA): pass per-request prompt + # lengths so the attention backend can keep the prefix globally visible. + # The backend owns the persistent CUDA-graph-safe GPU buffer. + rswa_prefix_lens = None + if self.model_config.rswa_window is not None: + rswa_prefix_lens = num_prompt_tokens_cpu + cm_base = CommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -2352,6 +2359,7 @@ class GPUModelRunner( is_prefilling=is_prefilling, positions=self.positions[:num_tokens_padded], mm_req_doc_ranges=req_doc_ranges, + rswa_prefix_lens=rswa_prefix_lens, ) if self.dcp_world_size > 1: From c7ca0bccae667934c29c654544131cdab046adfd Mon Sep 17 00:00:00 2001 From: Olga Miroshnichenko Date: Sun, 28 Jun 2026 10:04:08 +0300 Subject: [PATCH 084/138] [ROCm][Perf] Add Fused Shared Expert (FSE) support for GLM-4.5/6/7 (#44313) Signed-off-by: Olga Miroshnichenko Signed-off-by: Mehdi Ghanimifard Co-authored-by: Mehdi Ghanimifard Co-authored-by: Mehdi Ghanimifard --- vllm/model_executor/models/glm4_moe.py | 191 ++++++++++++++------- vllm/model_executor/models/glm4_moe_mtp.py | 154 ++++++++++++----- 2 files changed, 247 insertions(+), 98 deletions(-) diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 8226b65c45c..e3f94c673f4 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -32,6 +32,7 @@ import torch from torch import nn from transformers.models.glm4_moe import Glm4MoeConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( @@ -168,7 +169,16 @@ class Glm4MoE(nn.Module): self.physical_expert_start + self.n_local_physical_experts ) - if config.n_shared_experts is not None: + # AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py + # pattern (see Glm4MoE / FusedMoE wiring there). + self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + self.is_fusion_moe_shared_experts_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + + if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: + self.shared_experts = None + else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts self.shared_experts = Glm4MoeMLP( hidden_size=config.hidden_size, @@ -178,8 +188,6 @@ class Glm4MoE(nn.Module): reduce_results=False, prefix=f"{prefix}.shared_experts", ) - else: - self.shared_experts = None self.experts = FusedMoE( shared_experts=self.shared_experts, @@ -194,12 +202,18 @@ class Glm4MoE(nn.Module): topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func="sigmoid", + # aiter applies routed_scaling_factor internally; see deepseek_v2.py. routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_output=True, + 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, router_logits_dtype=torch.float32, + n_shared_experts=( + config.n_shared_experts + if self.is_fusion_moe_shared_experts_enabled + else None + ), ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -469,15 +483,25 @@ class Glm4MoeModel(nn.Module): def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + # FSE widens the mapping by n_shared_experts slots; see deepseek_v2.py. + num_experts = self.config.n_routed_experts + if ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + and self.config.n_shared_experts + ): + num_experts += self.config.n_shared_experts return fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts, + num_experts=num_experts, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + rocm_aiter_moe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -494,6 +518,11 @@ class Glm4MoeModel(nn.Module): spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is not None: continue + + is_fusion_moe_shared_experts_layer = ( + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + ) + 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: @@ -506,6 +535,8 @@ class Glm4MoeModel(nn.Module): # for mlp.experts[0].gate_gate_up_proj, which breaks load. if ("mlp.experts." in name) and name not in params_dict: continue + if is_fusion_moe_shared_experts_layer: + continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. @@ -527,65 +558,109 @@ class Glm4MoeModel(nn.Module): break else: is_expert_weight = False - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - # Anyway, this is an expert weight and should not be - # attempted to load as other weights later - is_expert_weight = True - - # Do not modify `name` since the loop may continue here - # Instead, create a new variable - name_mapped = name.replace(weight_name, param_name) - - if is_pp_missing_parameter(name_mapped, self): - continue - - param = params_dict[name_mapped] - # We should ask the weight loader to return success or not - # here since otherwise we may skip experts with other - # available replicas. - weight_loader = typing.cast( - Callable[..., bool], param.weight_loader + # FSE: split a widened mlp.shared_experts tensor into + # n_shared_experts chunks; see deepseek_v2.py for details. + num_chunks = 1 + split_dim = 0 + chunk_size = 0 + if is_fusion_moe_shared_experts_layer: + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1 + split_dim = ( + 1 + if ("down_proj.weight" in name and loaded_weight.ndim > 1) + else 0 ) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - break - else: - if is_expert_weight: - # We've checked that this is an expert weight - # However it's not mapped locally to this rank - # So we simply skip it - continue + total = loaded_weight.shape[split_dim] + if total % num_chunks != 0: + raise ValueError( + f"FSE shared-expert weight {name} has dim " + f"{total} along axis {split_dim} which is not " + f"divisible by n_shared_experts={num_chunks}." + ) + chunk_size = total // num_chunks - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue + for j in range(num_chunks): + chunk_name = name + weight_to_load = loaded_weight - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue + if is_fusion_moe_shared_experts_layer: + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size) + if loaded_weight.ndim == 1: + weight_to_load = loaded_weight[chunk_slice] + elif split_dim == 0: + weight_to_load = loaded_weight[chunk_slice, :] + else: + weight_to_load = loaded_weight[:, chunk_slice] + # Synthesize an expert-style name for expert mapping. + chunk_name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts + j}", + ) - if is_pp_missing_parameter(name, self): - continue + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in chunk_name: + continue - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) + # Anyway, this is an expert weight and should not be + # attempted to load as other weights later + is_expert_weight = True + + # Do not modify `name` since the loop may continue here + # Instead, create a new variable + name_mapped = chunk_name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name_mapped, self): + continue + + param = params_dict[name_mapped] + # We should ask the weight loader to return success + # or not here since otherwise we may skip experts + # with other available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + weight_to_load, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + if not is_fusion_moe_shared_experts_layer: + name = name_mapped + else: + loaded_params.add(name_mapped) + break + else: + if is_expert_weight: + # We've checked that this is an expert weight + # However it's not mapped locally to this rank + # So we simply skip it + continue + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + if name is not None and not is_fusion_moe_shared_experts_layer: + loaded_params.add(name) return loaded_params diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index b255b67d885..4d7b291df12 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -24,12 +24,14 @@ """Inference-only GLM-4.5, GLM-4.6, GLM-4.7 MTP model compatible with HuggingFace weights.""" -from collections.abc import Iterable +import typing +from collections.abc import Callable, Iterable import torch import torch.nn as nn from transformers import PretrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import CacheConfig, ParallelConfig, VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, @@ -237,6 +239,10 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts): return self.model.compute_logits(hidden_states, spec_step_idx) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # FSE weight loading mirrors glm4_moe.py / deepseek_mtp.py. + rocm_aiter_moe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -248,12 +254,15 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts): # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + num_experts = self.config.n_routed_experts + if rocm_aiter_moe_shared_expert_enabled and self.config.n_shared_experts: + num_experts += self.config.n_shared_experts expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts, + num_experts=num_experts, ) params_dict = dict(self.named_parameters()) @@ -269,6 +278,11 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts): if spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) + + is_fusion_moe_shared_experts_layer = ( + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + ) + 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: @@ -281,6 +295,8 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts): # for mlp.experts[0].gate_gate_up_proj, which breaks load. if ("mlp.experts." in name) and name not in params_dict: continue + if is_fusion_moe_shared_experts_layer: + continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: @@ -291,47 +307,105 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts): weight_loader(param, loaded_weight, shard_id) break else: - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader( - param, - loaded_weight, - name, - shard_id=shard_id, - expert_id=expert_id, + # FSE: split a widened mlp.shared_experts tensor into + # n_shared_experts chunks; see deepseek_v2.py for details. + num_chunks = 1 + split_dim = 0 + chunk_size = 0 + if is_fusion_moe_shared_experts_layer: + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1 + split_dim = ( + 1 + if ("down_proj.weight" in name and loaded_weight.ndim > 1) + else 0 ) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Some checkpoints include weight scale tensors for the - # LM head even when the quantized head isn't built. Skip - # them if the model does not expose a matching parameter - # to avoid KeyError during load. - if name.endswith(".weight_scale") and name not in params_dict: - continue + total = loaded_weight.shape[split_dim] + if total % num_chunks != 0: + raise ValueError( + f"FSE shared-expert weight {name} has dim " + f"{total} along axis {split_dim} which is " + f"not divisible by " + f"n_shared_experts={num_chunks}." + ) + chunk_size = total // num_chunks - # According to DeepSeek-V3 Technical Report, MTP modules - # shares embedding layer. We only load the first weights. - if ( - spec_layer != self.model.mtp_start_layer_idx - and ".layers" not in name - ): - continue + for j in range(num_chunks): + chunk_name = name + weight_to_load = loaded_weight - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) + if is_fusion_moe_shared_experts_layer: + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size) + if loaded_weight.ndim == 1: + weight_to_load = loaded_weight[chunk_slice] + elif split_dim == 0: + weight_to_load = loaded_weight[chunk_slice, :] + else: + weight_to_load = loaded_weight[:, chunk_slice] + chunk_name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts + j}", + ) + + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in chunk_name: + continue + + is_expert_weight = True + name_mapped = chunk_name.replace(weight_name, param_name) + + param = params_dict[name_mapped] + # Use return_success so we don't blindly mark + # remote-expert replicas as loaded on this rank. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + weight_to_load, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + if not is_fusion_moe_shared_experts_layer: + name = name_mapped + else: + loaded_params.add(name_mapped) + break + else: + if is_expert_weight: + # Expert weight not local to this rank; skip. + continue + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Some checkpoints include weight scale tensors for + # the LM head even when the quantized head isn't + # built. Skip them if the model does not expose a + # matching parameter to avoid KeyError during load. + if name.endswith(".weight_scale") and name not in params_dict: + continue + + # According to DeepSeek-V3 Technical Report, MTP + # modules share the embedding layer. We only load + # the first weights. + if ( + spec_layer != self.model.mtp_start_layer_idx + and ".layers" not in name + ): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + if not is_fusion_moe_shared_experts_layer: + loaded_params.add(name) return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: From 35e6c86caaaced4fd1398739fb04140b65d6ca89 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Sun, 28 Jun 2026 15:06:43 +0800 Subject: [PATCH 085/138] [Bugfix][MM][CG] Enable dual-path ViT CUDA graph for Step3-VL (#46034) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- docs/design/cuda_graphs_multimodal.md | 2 +- vllm/model_executor/models/step3_vl.py | 210 ++++++++++--------------- 2 files changed, 87 insertions(+), 125 deletions(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 264b1f139f6..ceefc195021 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -136,7 +136,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | โœ…๏ธŽ | โœ…๏ธŽ | โŒ๏ธŽ | | `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | โœ…๏ธŽ | โœ…๏ธŽ | โŒ๏ธŽ | | `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | โœ…๏ธŽ | โœ…๏ธŽ | โŒ๏ธŽ | -| `Step3VLForConditionalGeneration` | `Step3-VL` | โœ…๏ธŽ | โŒ๏ธŽ | โŒ๏ธŽ | +| `Step3VLForConditionalGeneration` | `Step3-VL` | โœ…๏ธŽ | โŒ๏ธŽ | โœ…๏ธŽ | !!! note Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 9e3cfbcff25..7b3bb93ad11 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -589,6 +589,31 @@ class Step3VLForConditionalGeneration( h2 = (h1 - 1) // 2 + 1 return h2 * h2 + @property + def img_output_tokens(self) -> int: + return self._compute_spatial_tokens( + self.config.vision_config.image_size, + self.config.vision_config.patch_size, + self.config.understand_projector_stride, + ) + + @property + def patch_output_tokens(self) -> int: + return self._compute_spatial_tokens( + 504, + self.config.vision_config.patch_size, + self.config.understand_projector_stride, + ) + + def _batched_encoder_forward( + self, + pixel_values: torch.Tensor, + ) -> torch.Tensor: + image_features = self._process_image_features( + self._get_vision_model_output(pixel_values) + ) + return image_features.reshape(-1, image_features.shape[-1]) + def _parse_and_validate_image_input( self, **kwargs: object ) -> Step3VLImageInputs | None: @@ -695,6 +720,8 @@ class Step3VLForConditionalGeneration( is_multimodal=is_multimodal, ) + # -- SupportsEncoderCudaGraph protocol methods -- + def get_encoder_cudagraph_config(self): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphConfig, @@ -707,18 +734,16 @@ class Step3VLForConditionalGeneration( "patch_pixel_values", ], out_hidden_size=self.config.hidden_size, + enable_dual_path_graph=True, + global_token_per_image=self.img_output_tokens, + local_token_per_patch=self.patch_output_tokens, ) def get_encoder_cudagraph_budget_range( self, vllm_config: "VllmConfig", ) -> tuple[int, int]: - # An image without patches - min_budget = self._compute_spatial_tokens( - self.config.vision_config.image_size, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) + min_budget = self.img_output_tokens max_budget = min( vllm_config.scheduler_config.max_num_batched_tokens, self.model_config.max_model_len, @@ -732,22 +757,6 @@ class Step3VLForConditionalGeneration( from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec num_patches = mm_kwargs.get("num_patches") - img_output_tokens = self._compute_spatial_tokens( - self.config.vision_config.image_size, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) - - # NOTE: 504 is the hard coded size for each patch after processing - # by the vision model, which is determined by the current architecture - # of the vision model and may need to be updated if the architecture changes. - # The number of tokens for each patch is calculated based on this - # size and the patch size. - patch_output_tokens = self._compute_spatial_tokens( - 504, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) img_grid = ( self.config.vision_config.image_size // self.config.vision_config.patch_size @@ -759,7 +768,11 @@ class Step3VLForConditionalGeneration( return [ EncoderItemSpec( input_size=(total_image_pixel + num_patch * total_patch_pixel), - output_tokens=(img_output_tokens + num_patch * patch_output_tokens), + output_tokens=( + self.img_output_tokens + num_patch * self.patch_output_tokens + ), + global_output_tokens=self.img_output_tokens, + local_output_tokens=num_patch * self.patch_output_tokens, ) for num_patch in num_patches ] @@ -810,46 +823,30 @@ class Step3VLForConditionalGeneration( EncoderCudaGraphCaptureInputs, ) - # For pixel_value, the max input size is max_batch_size - img_output_tokens = self._compute_spatial_tokens( - self.config.vision_config.image_size, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) - patch_output_tokens = self._compute_spatial_tokens( - 504, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) - dummy_pixel_values = torch.randn( - max_batch_size, - 3, - self.config.vision_config.image_size, - self.config.vision_config.image_size, - device=device, - dtype=dtype, - ) - # max_num_patches is the max total patches across the whole batch. - # token_budget = max_batch_size * img_out + max_num_patches * patch_out - max_num_patches = max( - 0, - (token_budget - max_batch_size * img_output_tokens) // patch_output_tokens, - ) - dummy_patch_pixel_values = torch.randn( - max_num_patches, - 3, - 504, - 504, - device=device, - dtype=dtype, - ) - # num_patches is NOT in values -- the per-item merge is done - # CPU-side by finalize_encoder_cudagraph_output using the actual - # batch's num_patches from mm_kwargs. - values = { - "pixel_values": dummy_pixel_values, - "patch_pixel_values": dummy_patch_pixel_values, - } + assert path in ("global", "local") + if path == "global": + max_num_images = token_budget // self.img_output_tokens + max_batch_size = min(max_batch_size, max_num_images) + dummy_pixel_values = torch.randn( + max_batch_size, + 3, + self.config.vision_config.image_size, + self.config.vision_config.image_size, + device=device, + dtype=dtype, + ) + values = {"pixel_values": dummy_pixel_values} + else: + max_num_patches = token_budget // self.patch_output_tokens + dummy_patch_pixel_values = torch.randn( + max_num_patches, + 3, + 504, + 504, + device=device, + dtype=dtype, + ) + values = {"patch_pixel_values": dummy_patch_pixel_values} return EncoderCudaGraphCaptureInputs( values=values, @@ -860,42 +857,22 @@ class Step3VLForConditionalGeneration( values: dict[str, torch.Tensor], path: str = "default", ) -> torch.Tensor: - # Graph captures only the compute (vision model + conv projector). - # Per-item merge happens CPU-side in finalize_encoder_cudagraph_output - # using actual num_patches from the batch data. - pixel_values = values["pixel_values"] - patch_pixel_values = values["patch_pixel_values"] - - image_features = self._process_image_features( - self._get_vision_model_output(pixel_values) - ) - - has_patches = len(patch_pixel_values) > 0 - if has_patches: - patch_features = self._process_image_features( - self._get_vision_model_output(patch_pixel_values) - ) - - # Deterministic single cat: [all_img_flat, all_patch_flat] - img_flat = image_features.reshape(-1, image_features.shape[-1]) - if has_patches: - patch_flat = patch_features.reshape(-1, patch_features.shape[-1]) - return torch.cat([img_flat, patch_flat], dim=0) - return img_flat + assert path in ("global", "local") + if path == "global": + return self._batched_encoder_forward(values["pixel_values"]) + else: + return self._batched_encoder_forward(values["patch_pixel_values"]) def encoder_eager_forward( self, mm_kwargs: dict[str, Any], path: str = "default", ) -> torch.Tensor: - image_input = Step3VLImagePixelInputs( - type="pixel_values", - pixel_values=mm_kwargs["pixel_values"], - patch_pixel_values=mm_kwargs["patch_pixel_values"], - num_patches=mm_kwargs["num_patches"], - ) - vision_embeddings = self._process_image_input(image_input) - return torch.cat(vision_embeddings, dim=0) + assert path in ("global", "local") + if path == "global": + return self._batched_encoder_forward(mm_kwargs["pixel_values"]) + else: + return self._batched_encoder_forward(mm_kwargs["patch_pixel_values"]) def postprocess_encoder_output( self, @@ -907,38 +884,24 @@ class Step3VLForConditionalGeneration( batch_mm_kwargs: dict[str, Any] | None = None, local_output: torch.Tensor | None = None, ): - """CPU-side per-item merge after graph replay. + """CPU-side per-item merge after dual-path graph replay. - The graph output is ``[all_img_flat, all_patch_flat]``. - This method splits the flat output into image and patch features, - then reassembles per-item embeddings using the *actual* batch - ``num_patches`` from ``batch_mm_kwargs`` (not the capture-time values). + ``output`` contains global-image features and ``local_output`` + contains local-patch features (or ``None`` when there are no patches). """ num_patches = batch_mm_kwargs["num_patches"] hidden = output.shape[-1] bsz = len(indices) - img_out = self._compute_spatial_tokens( - self.config.vision_config.image_size, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) - patch_out = self._compute_spatial_tokens( - 504, - self.config.vision_config.patch_size, - self.config.understand_projector_stride, - ) - - # Valid portion: bsz images, actual_total_patches patches actual_np = [int(np) for np in num_patches] total_patches = sum(actual_np) - img_tokens = bsz * img_out - patch_tokens = total_patches * patch_out + img_tokens = bsz * self.img_output_tokens + patch_tokens = total_patches * self.patch_output_tokens - img_part = output[:img_tokens].reshape(bsz, img_out, hidden) + global_part = output[:img_tokens].reshape(bsz, self.img_output_tokens, hidden) if total_patches > 0: - patch_part = output[img_tokens : img_tokens + patch_tokens].reshape( - -1, patch_out, hidden + patch_part = local_output[:patch_tokens].reshape( + -1, self.patch_output_tokens, hidden ) else: patch_part = None @@ -951,7 +914,7 @@ class Step3VLForConditionalGeneration( if patch_part is not None and np > 0: parts.append(patch_part[cur_patch : cur_patch + np].reshape(-1, hidden)) cur_patch += np - parts.append(img_part[i].reshape(-1, hidden)) + parts.append(global_part[i].reshape(-1, hidden)) merged[idx] = torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] out = [merged[i] for i in indices] @@ -969,14 +932,13 @@ class Step3VLForConditionalGeneration( EncoderCudaGraphReplayBuffers, ) - # Only patch_pixel_values lives in the values dict; num_patches is - # processed CPU-side by finalize_encoder_cudagraph_output. - return EncoderCudaGraphReplayBuffers( - values={ - "pixel_values": mm_kwargs["pixel_values"], - "patch_pixel_values": mm_kwargs["patch_pixel_values"], - }, - ) + assert path in ("global", "local") + if path == "global": + values = {"pixel_values": mm_kwargs["pixel_values"]} + else: + values = {"patch_pixel_values": mm_kwargs["patch_pixel_values"]} + + return EncoderCudaGraphReplayBuffers(values=values) def forward( self, From a2a92cbbaac1175de96fc6f4712b4bba789a0c02 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sun, 28 Jun 2026 02:07:14 -0500 Subject: [PATCH 086/138] [Hardware][AMD][CI] Tweak mirrored tests; improve CI base dependency change detection (#46930) Signed-off-by: Matthew Wong --- .buildkite/scripts/ci-bake-rocm.sh | 2 +- .buildkite/test-amd.yaml | 33 +++++++------------- .buildkite/test_areas/basic_correctness.yaml | 2 +- .buildkite/test_areas/distributed.yaml | 12 ------- .buildkite/test_areas/misc.yaml | 10 ++++++ docker/Dockerfile.rocm | 6 ++++ 6 files changed, 30 insertions(+), 35 deletions(-) diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 51cffb8e20d..4ccbbb352d9 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -15,7 +15,7 @@ set -euo pipefail DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" -DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base tools/install_torchcodec_rocm.sh tests/vllm_test_utils" +DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="1" diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a9608cc332f..7521901c9a6 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -438,7 +438,7 @@ steps: #----------------------------------------------------- mi300 ยท basic_correctness -----------------------------------------------------# - label: Basic Correctness # TBD - timeout_in_minutes: 40 + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 fast_check: true @@ -456,7 +456,7 @@ steps: - pytest -v -s basic_correctness/test_cpu_offload.py - label: Distributed Model Tests (2 GPUs) # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 num_gpus: 2 @@ -678,7 +678,7 @@ steps: - pytest -v -s distributed/test_eplb_spec_decode.py - label: Distributed Tests (2xH100-2xMI300) # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 num_gpus: 2 @@ -1199,7 +1199,7 @@ steps: #--------------------------------------------------------- mi300 ยท examples ----------------------------------------------------------# - label: Examples # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 optional: true @@ -1212,7 +1212,7 @@ steps: commands: - pip install tensorizer # Basic - - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN + - python3 basic/offline_inference/chat.py - python3 basic/offline_inference/generate.py --model facebook/opt-125m - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 - python3 basic/offline_inference/classify.py @@ -1220,11 +1220,8 @@ steps: - python3 basic/offline_inference/score.py # Multi-modal models - python3 generate/multimodal/audio_language_offline.py --seed 0 - # These two examples import transformers before vllm, so on ROCm the HIP context - # is initialized in the parent before vllm sets this guard, poisoning fork. Set it - # inline to keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). - - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_offline.py --seed 0 - - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 + - python3 generate/multimodal/vision_language_offline.py --seed 0 + - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 # Pooling models - python3 pooling/embed/vision_embedding_offline.py --seed 0 @@ -1666,10 +1663,7 @@ steps: - pytest -v -s tests/models/test_transformers.py - pytest -v -s tests/models/multimodal/test_mapping.py - python3 examples/basic/offline_inference/chat.py - # This example imports transformers before vllm, so on ROCm the HIP context is - # initialized in the parent before vllm sets this guard, poisoning fork. Set it - # inline to keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). - - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl + - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper #---------------------------------------------------------- mi300 ยท plugins ----------------------------------------------------------# @@ -2773,7 +2767,7 @@ steps: #--------------------------------------------------------- mi355 ยท examples ----------------------------------------------------------# - label: Examples # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 working_dir: "/vllm-workspace/examples" @@ -2785,7 +2779,7 @@ steps: commands: - pip install tensorizer # Basic - - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN + - python3 basic/offline_inference/chat.py - python3 basic/offline_inference/generate.py --model facebook/opt-125m - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 - python3 basic/offline_inference/classify.py @@ -2793,11 +2787,8 @@ steps: - python3 basic/offline_inference/score.py # Multi-modal models - python3 generate/multimodal/audio_language_offline.py --seed 0 - # These two examples import transformers before vllm, so on ROCm the HIP context - # is initialized in the parent before vllm sets this guard, poisoning fork. Set it - # inline to keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). - - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_offline.py --seed 0 - - PYTORCH_NVML_BASED_CUDA_CHECK=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 + - python3 generate/multimodal/vision_language_offline.py --seed 0 + - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0 - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0 # Pooling models - python3 pooling/embed/vision_embedding_offline.py --seed 0 diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index 7e166a8a28e..d7173b6438d 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -19,6 +19,6 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 40 + timeout_in_minutes: 50 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 5ff4b24b744..2cc52603c23 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -224,18 +224,6 @@ steps: - 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 - mirror: - amd: - device: mi300_2 - timeout_in_minutes: 180 - depends_on: - - image-build-amd - commands: - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 - - 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 (2xB200) key: distributed-tests-2xb200 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index f5db2e956b6..450cfcbc26d 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -224,6 +224,16 @@ steps: - python3 features/speculative_decoding/spec_decode_offline.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 # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU - python3 features/speculative_decoding/spec_decode_offline.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 + mirror: + amd: + device: mi325_1 + source_file_dependencies: + - vllm/entrypoints + - vllm/multimodal + - examples/ + - vllm/platforms/rocm.py + depends_on: + - image-build-amd - label: Metrics, Tracing (2 GPUs) key: metrics-tracing-2-gpus diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index dcae40c524a..c17444217f0 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -575,6 +575,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ENV HF_XET_HIGH_PERFORMANCE=1 ENV HF_HUB_DOWNLOAD_TIMEOUT=60 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 + # Pre-install vLLM test dependencies. COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt RUN --mount=type=cache,target=/root/.cache/uv \ @@ -695,6 +698,9 @@ ENV SAFETENSORS_FAST_GPU=1 # Performance environment variable. ENV HIP_FORCE_DEV_KERNARG=1 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 + # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" From 09841ae705ce73967b3303cdda5c7046d7710f5f Mon Sep 17 00:00:00 2001 From: Ranran Date: Sun, 28 Jun 2026 02:07:33 -0500 Subject: [PATCH 087/138] [Render][Speculator] Add return_loss_mask to render endpoint for training data generation (#46846) Signed-off-by: Ranran Haoran Zhang Co-authored-by: Benjamin Chislett --- tests/entrypoints/serve/render/test_render.py | 126 +++++++++++++++++- .../openai/chat_completion/protocol.py | 13 ++ vllm/entrypoints/serve/disagg/protocol.py | 8 ++ vllm/entrypoints/serve/render/serving.py | 22 +++ vllm/inputs/engine.py | 10 ++ vllm/renderers/hf.py | 125 ++++++++++++++--- vllm/renderers/params.py | 4 + 7 files changed, 288 insertions(+), 20 deletions(-) diff --git a/tests/entrypoints/serve/render/test_render.py b/tests/entrypoints/serve/render/test_render.py index d7339361ff7..ffd7f9f30ae 100644 --- a/tests/entrypoints/serve/render/test_render.py +++ b/tests/entrypoints/serve/render/test_render.py @@ -14,7 +14,7 @@ MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" @pytest.fixture(scope="module") def server(): - args: list[str] = [] + args: list[str] = ["--trust-request-chat-template"] with RemoteLaunchRenderServer(MODEL_NAME, args) as remote_server: yield remote_server @@ -369,3 +369,127 @@ async def test_completion_render_multiple_prompts_token_offsets(client): assert len(offsets) == len(item["token_ids"]) for start, end in offsets: assert 0 <= start <= end <= len(prompt) + + +@pytest.mark.asyncio +async def test_chat_completion_render_assistant_tokens_mask_default(client): + """Without return_assistant_tokens_mask, assistant_tokens_mask should be null.""" + response = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + {"role": "user", "content": "How are you?"}, + ], + }, + ) + + assert response.status_code == 200 + data = response.json() + assert data.get("assistant_tokens_mask") is None + + +@pytest.mark.asyncio +async def test_chat_completion_render_assistant_tokens_mask_false(client): + """Explicitly setting return_assistant_tokens_mask=false gives null.""" + response = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [ + {"role": "user", "content": "Hello"}, + ], + "return_assistant_tokens_mask": False, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert data.get("assistant_tokens_mask") is None + + +@pytest.mark.asyncio +async def test_chat_render_assistant_tokens_mask_null_without_gen_tags( + client, +): + """The tiny test model lacks ``{% generation %}`` tags, so the mask is null.""" + response = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + ], + "return_assistant_tokens_mask": True, + }, + ) + + assert response.status_code == 200 + assert response.json().get("assistant_tokens_mask") is None + + +# A minimal chat template with {% generation %} tags so we can test that +# the mask correctly marks assistant tokens. +_TEMPLATE_WITH_GENERATION = ( + "{% for m in messages %}" + "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n" + "{% elif m['role'] == 'assistant' %}" + "{% generation %}Assistant: {{ m['content'] }}\n{% endgeneration %}" + "{% endif %}" + "{% endfor %}" +) + + +@pytest.mark.asyncio +async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags( + client, +): + """With a ``{% generation %}``-enabled template, the mask marks assistant + tokens and the masked tokens decode to the assistant content.""" + response = await client.post( + "/v1/chat/completions/render", + json={ + "model": MODEL_NAME, + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + {"role": "user", "content": "Bye"}, + ], + "chat_template": _TEMPLATE_WITH_GENERATION, + "return_assistant_tokens_mask": True, + }, + ) + + assert response.status_code == 200 + data = response.json() + + mask = data["assistant_tokens_mask"] + token_ids = data["token_ids"] + assert mask is not None + assert isinstance(mask, list) + assert len(mask) == len(token_ids) + assert all(v in (0, 1) for v in mask) + assert sum(mask) > 0, "mask should mark at least one assistant token" + + # Detokenize masked (assistant) and unmasked (non-assistant) tokens + # separately to verify the mask is correct, not just non-empty. + masked_ids = [t for t, m in zip(token_ids, mask, strict=True) if m] + unmasked_ids = [t for t, m in zip(token_ids, mask, strict=True) if not m] + + detok = await client.post( + "/detokenize", + json={"model": MODEL_NAME, "tokens": masked_ids}, + ) + assert detok.status_code == 200 + assert "Hi!" in detok.json()["prompt"] + + detok_rest = await client.post( + "/detokenize", + json={"model": MODEL_NAME, "tokens": unmasked_ids}, + ) + assert detok_rest.status_code == 200 + assert "Hi!" not in detok_rest.json()["prompt"] + assert "Bye" in detok_rest.json()["prompt"] diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index aa2af69777c..36e467f6f32 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -407,6 +407,18 @@ class ChatCompletionRequest(OpenAIBaseModel): ), ) + return_assistant_tokens_mask: bool = Field( + default=False, + description=( + "If true, the /render response will include an " + "``assistant_tokens_mask`` field โ€” a per-token list of 0/1 " + "values indicating which tokens were assistant-generated. " + "Requires the chat template to use ``{% generation %}`` " + "tags. When the template does not support it, " + "``assistant_tokens_mask`` will be ``null``." + ), + ) + cache_salt: str | None = Field( default=None, description=( @@ -520,6 +532,7 @@ class ChatCompletionRequest(OpenAIBaseModel): extra_kwargs, ), media_io_kwargs=self.media_io_kwargs, + return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask), ) def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index d20752a9063..723c2792491 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -75,6 +75,14 @@ class GenerateRequest(BaseModel): token_ids: list[int] = Field(min_length=1) """The token ids to generate text from.""" + assistant_tokens_mask: list[int] | None = None + """Per-token mask (1 = assistant-generated, 0 = not). + + Only populated when the render request sets ``return_assistant_tokens_mask=True`` + and the chat template supports ``{% generation %}``. + ``None`` when the mask was not requested or could not be computed. + """ + @field_validator("token_ids") @classmethod def validate_token_ids(cls, v: list[int]) -> list[int]: diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 1bba26722b9..adcaf8af9af 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -127,11 +127,33 @@ class ServingRender(BaseServing): ) params = request.to_sampling_params(max_tokens, self.default_sampling_params) + assistant_tokens_mask: list[int] | None = engine_input.get( # type: ignore[assignment] + "assistant_tokens_mask" + ) + if assistant_tokens_mask is not None and len(assistant_tokens_mask) != len( + token_ids + ): + logger.warning( + "assistant_tokens_mask length (%d) != token_ids length (%d); " + "this can happen with multimodal inputs where " + "placeholder expansion changes the token count. " + "The mask may be positionally misaligned.", + len(assistant_tokens_mask), + len(token_ids), + ) + if len(assistant_tokens_mask) < len(token_ids): + assistant_tokens_mask.extend( + [0] * (len(token_ids) - len(assistant_tokens_mask)) + ) + else: + assistant_tokens_mask = assistant_tokens_mask[: len(token_ids)] + request_id = f"chatcmpl-{random_uuid()}" return GenerateRequest( request_id=request_id, token_ids=token_ids, + assistant_tokens_mask=assistant_tokens_mask, features=self._extract_mm_features(engine_input), sampling_params=params, model=request.model, diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index eacadcbc924..f997004d2fb 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -42,6 +42,11 @@ class TokensInput(_InputOptions): """Char-level (start, end) offsets per token, propagated from the renderer's TokensPrompt when offsets were computed.""" + assistant_tokens_mask: NotRequired[list[int] | None] + """Per-token 0/1 mask marking assistant-generated tokens. + Populated when ``return_assistant_tokens_mask=True`` is set on the + render request and the chat template supports ``{% generation %}``.""" + def tokens_input( prompt_token_ids: list[int], @@ -151,6 +156,11 @@ class MultiModalInput(_InputOptions): `prompt_token_ids`. """ + assistant_tokens_mask: NotRequired[list[int] | None] + """Per-token 0/1 mask marking assistant-generated tokens. + Populated when ``return_assistant_tokens_mask=True`` is set on the + render request and the chat template supports ``{% generation %}``.""" + def mm_input( prompt_token_ids: list[int], diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index ea0902c8806..490f589af45 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -7,7 +7,7 @@ import inspect import itertools import weakref from collections import defaultdict, deque -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -678,6 +678,7 @@ def safe_apply_chat_template( tools: list[dict[str, Any]] | None = ..., chat_template: str | None = ..., tokenize: Literal[True] = ..., + return_assistant_tokens_mask: Literal[False] = ..., **kwargs, ) -> list[int]: ... @overload @@ -689,8 +690,20 @@ def safe_apply_chat_template( tools: list[dict[str, Any]] | None = ..., chat_template: str | None = ..., tokenize: Literal[False] = ..., + return_assistant_tokens_mask: Literal[False] = ..., **kwargs, ) -> str: ... +@overload +def safe_apply_chat_template( + model_config: ModelConfig, + tokenizer: HfTokenizer, + conversation: list[ConversationMessage], + *, + tools: list[dict[str, Any]] | None = ..., + chat_template: str | None = ..., + return_assistant_tokens_mask: Literal[True], + **kwargs, +) -> tuple[list[int], list[int] | None]: ... def safe_apply_chat_template( model_config: ModelConfig, tokenizer: HfTokenizer, @@ -699,8 +712,9 @@ def safe_apply_chat_template( tools: list[dict[str, Any]] | None = None, chat_template: str | None = None, tokenize: bool = True, + return_assistant_tokens_mask: bool = False, **kwargs, -) -> str | list[int]: +) -> str | list[int] | tuple[list[int], list[int] | None]: chat_template = resolve_chat_template( tokenizer, chat_template=chat_template, @@ -728,6 +742,38 @@ def safe_apply_chat_template( chat_template_kwargs=kwargs, ) + # assistant_tokens_mask requires tokenized output โ€” force tokenize=True. + if return_assistant_tokens_mask: + tokenize = True + + # When return_assistant_tokens_mask is requested and the template supports it, + # request assistant_tokens_mask via return_dict. + # Check for the actual Jinja tag, not just the word "generation" + # (which also appears in add_generation_prompt). + if return_assistant_tokens_mask and "{% generation %}" in chat_template: + resolved_kwargs["return_assistant_tokens_mask"] = True + resolved_kwargs["return_dict"] = True + resolved_kwargs.pop("tokenize", None) + try: + result = tokenizer.apply_chat_template( + conversation=conversation, # type: ignore[arg-type] + tools=tools, # type: ignore[arg-type] + chat_template=chat_template, + tokenize=True, + **resolved_kwargs, + ) + except (TypeError, ValueError) as exc: + logger.warning( + "apply_chat_template failed for assistant_tokens_mask: %s", exc + ) + else: + if isinstance(result, Mapping): + token_ids = list(result.get("input_ids", [])) + mask_raw = result.get("assistant_masks") + mask = list(mask_raw) if mask_raw is not None else None + return token_ids, mask + return list(result), None + # transformers v5 changed the default of `return_dict` to True, which # makes `apply_chat_template(tokenize=True)` return a `BatchEncoding` # instead of `list[int]`. Force `return_dict=False` so downstream code @@ -737,23 +783,24 @@ def safe_apply_chat_template( resolved_kwargs["return_dict"] = False try: - return tokenizer.apply_chat_template( + plain = tokenizer.apply_chat_template( conversation=conversation, # type: ignore[arg-type] tools=tools, # type: ignore[arg-type] chat_template=chat_template, tokenize=tokenize, **resolved_kwargs, ) - # External library exceptions can sometimes occur despite the framework's - # internal exception management capabilities. except Exception as e: - # Log and report any library-related exceptions for further - # investigation. logger.exception( "An error occurred in `transformers` while applying chat template" ) raise ValueError(str(e)) from e + if return_assistant_tokens_mask: + assert isinstance(plain, list), f"Expected list[int], got {type(plain)}" + return plain, None + return plain + def rebuild_mm_uuids_from_mm_data( mm_uuids: MultiModalUUIDDict, @@ -934,12 +981,22 @@ class HfRenderer(BaseRenderer[HfTokenizer]): logger.warning_once(_TOKENIZE_OVERRIDE_WARNING) chat_template_kwargs["tokenize"] = True - prompt_raw = safe_apply_chat_template( - model_config, - tokenizer, - conversation, - **chat_template_kwargs, - ) + assistant_tokens_mask: list[int] | None = None + if params.return_assistant_tokens_mask: + prompt_raw, assistant_tokens_mask = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + return_assistant_tokens_mask=True, + **chat_template_kwargs, + ) + else: + prompt_raw = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + **chat_template_kwargs, + ) # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5 # model which uses unified vision chunks for both images and videos. @@ -965,6 +1022,9 @@ class HfRenderer(BaseRenderer[HfTokenizer]): prompt = parse_dec_only_prompt(prompt_raw) + if assistant_tokens_mask is not None: + cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask + # When `prompt_embeds` is mixed with other modality data, # `_process_tokens` runs `_process_multimodal` first (expanding # `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then @@ -1038,12 +1098,30 @@ class HfRenderer(BaseRenderer[HfTokenizer]): logger.warning_once(_TOKENIZE_OVERRIDE_WARNING) chat_template_kwargs["tokenize"] = True - prompt_raw = await self._apply_chat_template_async( - model_config, - tokenizer, - conversation, - **chat_template_kwargs, - ) + assistant_tokens_mask: list[int] | None = None + if params.return_assistant_tokens_mask: + result_with_mask = cast( + tuple[list[int], list[int] | None], + await make_async( + safe_apply_chat_template, + executor=self._executor, + )( + model_config, + tokenizer, + conversation, + return_assistant_tokens_mask=True, # type: ignore[arg-type] + **chat_template_kwargs, + ), + ) + prompt_raw: str | list[int] = result_with_mask[0] + assistant_tokens_mask = result_with_mask[1] + else: + prompt_raw = await self._apply_chat_template_async( + model_config, + tokenizer, + conversation, + **chat_template_kwargs, + ) # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5 # model which uses unified vision chunks for both images and videos. @@ -1067,6 +1145,9 @@ class HfRenderer(BaseRenderer[HfTokenizer]): prompt = parse_dec_only_prompt(prompt_raw) + if assistant_tokens_mask is not None: + cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask + # See `render_messages` for the rationale. if prompt_embeds_tensors and mm_data: assert prompt_embeds_placeholder_token_id is not None @@ -1108,6 +1189,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]): processor records all placeholder offsets in the final (post-expansion) coordinate space, no offset shifting needed afterwards. """ + assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None) prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) if prompt_embeds_info is not None: tensors, placeholder_token_id = prompt_embeds_info @@ -1123,6 +1205,8 @@ class HfRenderer(BaseRenderer[HfTokenizer]): tensors, mm_updates, ) + if assistant_tokens_mask is not None: + engine_input["assistant_tokens_mask"] = assistant_tokens_mask return engine_input @override @@ -1133,6 +1217,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]): skip_mm_cache: bool = False, ) -> TokensInput | MultiModalInput: """Async equivalent of `_process_tokens`.""" + assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None) prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) if prompt_embeds_info is not None: tensors, placeholder_token_id = prompt_embeds_info @@ -1150,6 +1235,8 @@ class HfRenderer(BaseRenderer[HfTokenizer]): tensors, mm_updates, ) + if assistant_tokens_mask is not None: + engine_input["assistant_tokens_mask"] = assistant_tokens_mask return engine_input @staticmethod diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 8e0aaf303cc..7e3670c738d 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -87,6 +87,9 @@ class ChatParams: mm_processor_kwargs: dict[str, Any] | None = None """The kwargs to pass to the multi-modal processor.""" + return_assistant_tokens_mask: bool = False + """Request a per-token assistant mask from apply_chat_template.""" + def with_defaults( self, default_chat_template_kwargs: dict[str, Any] | None = None, @@ -115,6 +118,7 @@ class ChatParams: default_mm_processor_kwargs, self.mm_processor_kwargs, ), + return_assistant_tokens_mask=self.return_assistant_tokens_mask, ) def get_apply_chat_template_kwargs(self) -> dict[str, Any]: From 6eb63a1da6996abad00323dc7e845dc868996524 Mon Sep 17 00:00:00 2001 From: frida-andersson Date: Sun, 28 Jun 2026 10:37:44 +0200 Subject: [PATCH 088/138] [Bugfix][DSv3.2] Skip indexer weights for index-cache-skipped layers (#46600) Signed-off-by: Frida Andersson Co-authored-by: Andreas Karatzas --- vllm/model_executor/models/deepseek_v2.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 09960050c06..aaca07b6930 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1451,6 +1451,11 @@ class DeepseekV2Model(nn.Module): pp_missing_layer_names = get_pp_missing_layer_names(self) params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + # With index_topk_freq>1 only some layers build an indexer, yet the + # checkpoint ships indexer weights for all of them; track the built ones. + indexer_present_prefixes = { + n.rsplit(".indexer.", 1)[0] for n in params_dict if ".indexer." in n + } for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue @@ -1459,6 +1464,11 @@ class DeepseekV2Model(nn.Module): if spec_layer is not None: continue # skip spec decode layers for main model + if ".indexer." in name and ( + name.rsplit(".indexer.", 1)[0] not in indexer_present_prefixes + ): + continue # this layer has no indexer; drop its checkpoint weights + is_fusion_moe_shared_experts_layer = ( rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) ) From 5ecae3266cd5e6b814e14368d615957ffe85fdef Mon Sep 17 00:00:00 2001 From: xaguilar-amd Date: Sun, 28 Jun 2026 16:52:00 +0200 Subject: [PATCH 089/138] [ROCm][Perf][MLA] Add AITER FlashAttention MLA prefill backend (`ROCM_AITER_FA`) (#45033) Signed-off-by: Xavier Aguilar Signed-off-by: Xavier Aguilar Co-authored-by: TJian --- .../v1/attention/test_mla_prefill_registry.py | 17 +++ .../v1/attention/test_mla_prefill_selector.py | 119 +++++++++++++++++ .../backends/mla/prefill/aiter_flash_attn.py | 121 ++++++++++++++++++ .../backends/mla/prefill/registry.py | 4 + .../backends/mla/prefill/selector.py | 8 ++ 5 files changed, 269 insertions(+) create mode 100644 vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index 668c17c3f55..dfa3a029cea 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/tests/v1/attention/test_mla_prefill_registry.py @@ -133,3 +133,20 @@ def test_clear_override(): def test_unknown_backend_name_raises(): with pytest.raises(ValueError, match="Unknown MLA prefill backend"): MLAPrefillBackendEnum["NONEXISTENT"] + + +def test_rocm_aiter_fa_registered(): + """ROCM_AITER_FA is a known backend pointing at the AITER FA class.""" + assert "ROCM_AITER_FA" in MLAPrefillBackendEnum.__members__ + + path = MLAPrefillBackendEnum.ROCM_AITER_FA.get_path() + assert path == ( + "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn." + "AiterFlashAttnPrefillBackend" + ) + + backend_cls = MLAPrefillBackendEnum.ROCM_AITER_FA.get_class() + assert backend_cls.get_name() == "ROCM_AITER_FA" + # The AITER FA path is the fp16/bf16 generic-varlen prefill path. + assert backend_cls.supports_dtype(torch.bfloat16) + assert backend_cls.supports_dtype(torch.float16) diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index 54e68e03f26..c8932032467 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -14,6 +14,7 @@ from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnu from vllm.v1.attention.backends.mla.prefill.selector import ( MLAPrefillSelectorConfig, _auto_select_mla_prefill_backend, + _get_mla_prefill_backend_priorities, get_mla_prefill_backend, ) @@ -166,6 +167,7 @@ class TestAutoSelectMLAPrefillBackend: return with ( + patch("vllm.platforms.current_platform") as mock_platform, patch.object( MLAPrefillBackendEnum.FLASH_ATTN, "get_class", @@ -173,6 +175,8 @@ class TestAutoSelectMLAPrefillBackend: ), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): + # Force the non-ROCm priority on the Blackwell. + mock_platform.is_rocm.return_value = False backend = _auto_select_mla_prefill_backend( capability, selector_config, @@ -272,6 +276,121 @@ class TestBackendValidation: assert invalid_reasons == [] +class TestROCmAiterFAPrefillSelection: + """Tests for the ROCm AITER FlashAttention MLA prefill backend.""" + + def test_rocm_priorities_prefer_aiter_fa(self): + """On ROCm, ROCM_AITER_FA is tried first, FLASH_ATTN as fallback.""" + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.is_rocm.return_value = True + priorities = _get_mla_prefill_backend_priorities( + DeviceCapability(major=9, minor=5) + ) + + assert priorities == [ + MLAPrefillBackendEnum.ROCM_AITER_FA, + MLAPrefillBackendEnum.FLASH_ATTN, + ] + + def test_supported_dtypes_are_fp16_bf16_only(self): + from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( + AiterFlashAttnPrefillBackend, + ) + + assert AiterFlashAttnPrefillBackend.supports_dtype(torch.bfloat16) + assert AiterFlashAttnPrefillBackend.supports_dtype(torch.float16) + # FP8 is served by the separate AITER ASM backend, not this one. + assert not AiterFlashAttnPrefillBackend.supports_dtype(torch.float8_e4m3fn) + + def test_supports_compute_capability_on_rocm(self): + from vllm.v1.attention.backends.mla.prefill import aiter_flash_attn as mod + + # Gating is decided by on_mi3xx(), not by capability + capability = MagicMock() + + with patch.object(mod.current_platform, "is_rocm", return_value=False): + assert not mod.AiterFlashAttnPrefillBackend.supports_compute_capability( + capability + ) + + with ( + patch.object(mod.current_platform, "is_rocm", return_value=True), + patch("vllm.platforms.rocm.on_mi3xx", return_value=False), + ): + assert not mod.AiterFlashAttnPrefillBackend.supports_compute_capability( + capability + ) + + with ( + patch.object(mod.current_platform, "is_rocm", return_value=True), + patch("vllm.platforms.rocm.on_mi3xx", return_value=True), + ): + assert mod.AiterFlashAttnPrefillBackend.supports_compute_capability( + capability + ) + + def test_is_available_delegates_to_rocm_aiter_ops(self): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla.prefill import aiter_flash_attn as mod + + with patch.object(rocm_aiter_ops, "is_enabled", return_value=False): + assert not mod.AiterFlashAttnPrefillBackend.is_available() + + with patch.object(rocm_aiter_ops, "is_enabled", return_value=True): + assert mod.AiterFlashAttnPrefillBackend.is_available() + + def test_auto_select_prefers_aiter_fa_on_rocm(self): + from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( + AiterFlashAttnPrefillBackend, + ) + + # gfx gating is simulated via the mocked validate_configuration, + # not the capability. + capability = MagicMock() + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) + + with ( + patch("vllm.platforms.current_platform") as mock_platform, + patch.object( + AiterFlashAttnPrefillBackend, + "validate_configuration", + return_value=[], + ), + ): + mock_platform.is_rocm.return_value = True + backend = _auto_select_mla_prefill_backend(capability, selector_config) + assert backend.get_name() == "ROCM_AITER_FA" + + def test_auto_select_falls_back_to_flash_attn_when_aiter_invalid(self): + from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( + AiterFlashAttnPrefillBackend, + ) + + try: + flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() + except ImportError: + pytest.skip("FLASH_ATTN backend not available") + return + + # the fallback is forced by the mocked validate_configuration, + # not the capability. + capability = MagicMock() + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) + + with ( + patch("vllm.platforms.current_platform") as mock_platform, + patch.object( + AiterFlashAttnPrefillBackend, + "validate_configuration", + return_value=["compute capability not supported"], + ), + patch.object(flash_attn_cls, "validate_configuration", return_value=[]), + ): + mock_platform.is_rocm.return_value = True + backend = _auto_select_mla_prefill_backend(capability, selector_config) + assert backend.get_name() == "FLASH_ATTN" + + class TestMLAPrefillBackendParsing: """Tests for string-based mla_prefill_backend parsing from CLI args.""" diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py new file mode 100644 index 00000000000..130fcf394be --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AITER FlashAttention backend for MLA prefill (ROCm). + +This backend calls ``aiter.flash_attn_varlen_func`` directly, which natively +supports different q/k and v head dims (qk headdim 192, v headdim 128) without +padding V, and dispatches to the fast ``aiter::fmha_fwd_`` kernel on +gfx942/gfx950 (fp16/bf16). +""" + +from typing import TYPE_CHECKING + +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.platforms.interface import DeviceCapability + + +class AiterFlashAttnPrefillBackend(MLAPrefillBackend): + """AITER FlashAttention backend for MLA prefill""" + + @staticmethod + def get_name() -> str: + return "ROCM_AITER_FA" + + @classmethod + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_mi3xx + + return on_mi3xx() + + @classmethod + def is_available(cls) -> bool: + from vllm._aiter_ops import rocm_aiter_ops + + return rocm_aiter_ops.is_enabled() + + def __init__( + self, + num_heads: int, + scale: float, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + vllm_config: "VllmConfig", + ) -> None: + super().__init__( + num_heads=num_heads, + scale=scale, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + vllm_config=vllm_config, + ) + + from aiter import flash_attn_varlen_func + + self.flash_attn_varlen_func = flash_attn_varlen_func + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + out: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert output_scale is None, ( + "AiterFlashAttnPrefillBackend does not support fused quantized output." + ) + result = self.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=self._prefill_metadata.query_start_loc, + cu_seqlens_k=self._prefill_metadata.query_start_loc, + max_seqlen_q=self._prefill_metadata.max_query_len, + max_seqlen_k=self._prefill_metadata.max_query_len, + softmax_scale=self.scale, + causal=True, + return_lse=return_softmax_lse, + out=out, + ) + + # aiter returns the bare output tensor when return_lse is False, and + # (out, softmax_lse) when it is True. + if return_softmax_lse: + return result[0], result[1] + return result + + def run_prefill_context_chunk( + self, + chunk_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self._prefill_metadata.chunked_context is not None + chunked = self._prefill_metadata.chunked_context + out, lse = self.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=self._prefill_metadata.query_start_loc, + cu_seqlens_k=chunked.cu_seq_lens[chunk_idx], + max_seqlen_q=self._prefill_metadata.max_query_len, + max_seqlen_k=chunked.max_seq_lens[chunk_idx], + softmax_scale=self.scale, + causal=False, + return_lse=True, + ) + return out, lse diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index 9c83ea1b13d..0d818a084ba 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -48,6 +48,10 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla." "TokenspeedMLAPrefillBackend" ) + ROCM_AITER_FA = ( + "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn." + "AiterFlashAttnPrefillBackend" + ) # 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/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index e100c098acb..a38b274dfcb 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -56,6 +56,14 @@ def _get_mla_prefill_backend_priorities( Returns: List of backends in priority order (highest priority first). """ + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + return [ + MLAPrefillBackendEnum.ROCM_AITER_FA, + MLAPrefillBackendEnum.FLASH_ATTN, + ] + if device_capability.major == 10: # Blackwell return [ MLAPrefillBackendEnum.FLASH_ATTN, From 5c91039c41bc0b6a4a4ab2dc5f62115946e38a30 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sun, 28 Jun 2026 10:55:54 -0400 Subject: [PATCH 090/138] [GLM5.2 Perf] Replace MOE all-reduce with reduce-scatter, 3.1%~3.2 E2E Throughput improvement (#46635) Signed-off-by: yewentao256 --- vllm/model_executor/models/deepseek_v2.py | 98 ++++++++++++++++++++--- 1 file changed, 86 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index aaca07b6930..2c6b075ae74 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -42,6 +42,7 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul @@ -128,6 +129,7 @@ class DeepseekAttention(nn.Module): max_position_embeddings: int = 8192, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, prefix: str = "", **kwargs, ) -> None: @@ -166,6 +168,7 @@ class DeepseekAttention(nn.Module): self.total_num_heads * self.head_dim, hidden_size, bias=False, + reduce_results=reduce_results, quant_config=quant_config, ) @@ -372,15 +375,17 @@ class DeepseekV2MoE(nn.Module): self.gate.e_score_correction_bias.data.to(self.gate.out_dtype) ) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward( + self, + hidden_states: torch.Tensor, + already_sequence_parallel: bool = False, + ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) # Chunk the hidden states so they aren't replicated across TP ranks. # This avoids duplicate computation in self.experts. - # TODO: We can replace the all_reduce at the end of attn with a - # reduce_scatter instead of chunking here. - if self.is_sequence_parallel: + if self.is_sequence_parallel and not already_sequence_parallel: hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: @@ -393,7 +398,7 @@ class DeepseekV2MoE(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - if self.is_sequence_parallel: + if self.is_sequence_parallel and not already_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) @@ -436,6 +441,7 @@ class DeepseekV2Attention(nn.Module): cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, topk_indices_buffer: torch.Tensor | None = None, + reduce_results: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -502,6 +508,7 @@ class DeepseekV2Attention(nn.Module): self.num_heads * self.v_head_dim, self.hidden_size, bias=False, + reduce_results=reduce_results, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -950,6 +957,7 @@ class DeepseekV2MLAAttention(nn.Module): prefix: str = "", topk_indices_buffer: torch.Tensor | None = None, input_size: int | None = None, + reduce_results: bool = True, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -1018,6 +1026,7 @@ class DeepseekV2MLAAttention(nn.Module): self.num_heads * self.v_head_dim, self.hidden_size, bias=False, + reduce_results=reduce_results, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -1184,6 +1193,18 @@ class DeepseekV2DecoderLayer(nn.Module): attn_cls = DeepseekV2MLAAttention else: attn_cls = DeepseekV2Attention + is_moe_layer = ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % moe_layer_freq == 0 + ) + # TODO(wentao): enable SP MoE with PP after the PP boundary logic can safely + # send/receive sequence-parallel hidden_states across stages. + self.use_sequence_parallel_moe = ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + and is_moe_layer + ) self.self_attn = attn_cls( vllm_config=vllm_config, config=config, @@ -1199,13 +1220,10 @@ class DeepseekV2DecoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, + reduce_results=not self.use_sequence_parallel_moe, ) - if ( - config.n_routed_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % moe_layer_freq == 0 - ): + if is_moe_layer: self.mlp = DeepseekV2MoE( config=config, parallel_config=parallel_config, @@ -1233,6 +1251,13 @@ class DeepseekV2DecoderLayer(nn.Module): residual: torch.Tensor | None, llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: + full_num_tokens = positions.shape[0] + input_is_sequence_parallel = ( + self.use_sequence_parallel_moe + and residual is not None + and hidden_states.shape[0] != full_num_tokens + ) + # Self Attention if residual is None: residual = hidden_states @@ -1240,6 +1265,10 @@ class DeepseekV2DecoderLayer(nn.Module): else: hidden_states, residual = self.input_layernorm(hidden_states, residual) + if input_is_sequence_parallel: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + attn_kwargs = { "positions": positions, "hidden_states": hidden_states, @@ -1261,9 +1290,29 @@ class DeepseekV2DecoderLayer(nn.Module): # first layer. residual *= 1.0 / self.routed_scaling_factor + if self.use_sequence_parallel_moe: + sp_remainder = ( + hidden_states.shape[0] % get_tensor_model_parallel_world_size() + ) + # pad if not divisible by world size + if sp_remainder: + sp_pad = get_tensor_model_parallel_world_size() - sp_remainder + hidden_states = torch.nn.functional.pad( + hidden_states, (0, 0, 0, sp_pad) + ) + hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) + if not input_is_sequence_parallel: + residual = sequence_parallel_chunk(residual) + # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) + if self.use_sequence_parallel_moe: + hidden_states = self.mlp( + hidden_states, + already_sequence_parallel=True, + ) + else: + hidden_states = self.mlp(hidden_states) if isinstance(self.mlp, DeepseekV2MLP) and hidden_states.dtype == torch.float16: # Fix FP16 overflow @@ -1385,8 +1434,25 @@ class DeepseekV2Model(nn.Module): islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): + # all gather if we need to use the whole states + if ( + hidden_states.shape[0] != positions.shape[0] + and not layer.use_sequence_parallel_moe + ): + combined_states = torch.cat([hidden_states, residual], dim=-1) + combined_states = tensor_model_parallel_all_gather(combined_states, 0) + combined_states = combined_states[: positions.shape[0]] + hidden_states, residual = combined_states.split( + [self.hidden_size, self.hidden_size], dim=-1 + ) if idx in self.aux_hidden_state_layers: - aux_hidden_states.append(hidden_states + residual) + aux_hidden_state = hidden_states + residual + if aux_hidden_state.shape[0] != positions.shape[0]: + aux_hidden_state = tensor_model_parallel_all_gather( + aux_hidden_state, 0 + ) + aux_hidden_state = aux_hidden_state[: positions.shape[0]] + aux_hidden_states.append(aux_hidden_state) hidden_states, residual = layer( positions, hidden_states, residual, llama_4_scaling ) @@ -1396,6 +1462,14 @@ class DeepseekV2Model(nn.Module): {"hidden_states": hidden_states, "residual": residual} ) + if hidden_states.shape[0] != positions.shape[0]: + combined_states = torch.cat([hidden_states, residual], dim=-1) + combined_states = tensor_model_parallel_all_gather(combined_states, 0) + combined_states = combined_states[: positions.shape[0]] + hidden_states, residual = combined_states.split( + [self.hidden_size, self.hidden_size], dim=-1 + ) + hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states From 89876b0c548afdd932d41d5ef81b14347edb9ff7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 28 Jun 2026 08:17:39 -0700 Subject: [PATCH 091/138] [GLM5] Implement op fusion for GLM5/DSV3.2 (#46876) --- .../test_fused_deepseek_v32_norm_rope.py | 423 +++++++++ .../model_executor/layers/fused_moe/config.py | 6 + .../layers/fused_moe/runner/moe_runner.py | 2 + .../layers/sparse_attn_indexer.py | 10 +- vllm/models/deepseek_v32/nvidia/attention.py | 251 ++++-- vllm/models/deepseek_v32/nvidia/fused_ops.py | 63 ++ vllm/models/deepseek_v32/nvidia/kernels.py | 823 ++++++++++++++++++ vllm/models/deepseek_v32/nvidia/model.py | 24 +- vllm/models/deepseek_v32/nvidia/mtp.py | 21 +- 9 files changed, 1531 insertions(+), 92 deletions(-) create mode 100644 tests/kernels/test_fused_deepseek_v32_norm_rope.py create mode 100644 vllm/models/deepseek_v32/nvidia/fused_ops.py create mode 100644 vllm/models/deepseek_v32/nvidia/kernels.py diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py new file mode 100644 index 00000000000..a6f6d71b482 --- /dev/null +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the horizontally-fused deepseek_v32 (NVIDIA SM100) Triton +kernels used by the specialized DSA model: + + fused_norm_rope + - q : q_lora RMSNorm + - kv : kv_lora RMSNorm + (interleaved) RoPE on k_pe + MLA cache insert + (bf16 or per-tensor fp8) + - idx: indexer-K LayerNorm + RoPE (interleaved or NeoX) + UE8M0 fp8 quant + + packed indexer cache insert; plus the top-k buffer (-1) fill + fused_q + - mqa: ql_nope + (interleaved) RoPE'd q_pe, concat-quantized to the fp8 MQA + query + - idx: indexer-Q RoPE (interleaved or NeoX) + UE8M0 fp8 quant + folded + index weights + fused_eh_norm (MTP): zero-at-pos-0 + enorm RMSNorm(embeds) + hnorm + RMSNorm(prev), concatenated side-by-side + +Each kernel is compared against a PyTorch reference. The kernel keeps the whole +pipeline in fp32 and rounds once, so it can land on the opposite side of a +round-to-nearest tie from the reference for a few elements: deterministic fp8 +outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs use +rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses). +""" + +import pytest +import torch + +from vllm.models.deepseek_v32.nvidia import kernels as K +from vllm.platforms import current_platform + +FP8 = torch.float8_e4m3fn +FP8_MAX = 448.0 + +# GLM-5.2 / DeepSeek-V3.2 shapes (TP8 local heads). +Q_LORA = 2048 +KV_LORA = 512 +ROPE_DIM = 64 +NUM_HEADS = 8 +INDEX_HEADS = 32 +INDEX_HEAD_DIM = 128 +HIDDEN = 6144 +EPS = 1e-6 + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(89), + reason="deepseek_v32 fused kernels require CUDA with fp8 (SM89+)", +) + + +# โ”€โ”€ reference helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def make_cos_sin(max_pos: int, rot_dim: int, device) -> torch.Tensor: + """cos||sin cache: row[pos] = [cos(theta)(rot/2), sin(theta)(rot/2)].""" + half = rot_dim // 2 + inv_freq = 1.0 / ( + 10000.0 ** (torch.arange(0, half, dtype=torch.float32, device=device) / half) + ) + t = torch.arange(max_pos, dtype=torch.float32, device=device) + freqs = torch.einsum("i,j->ij", t, inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], dim=-1) + + +def rms_norm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """RMSNorm matching kernels._rms_norm (fp32, eps inside rsqrt). Returns fp32.""" + xf = x.float() + ms = xf.pow(2).mean(dim=-1, keepdim=True) + return xf * torch.rsqrt(ms + EPS) * w.float() + + +def layer_norm(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + xf = x.float() + mean = xf.mean(dim=-1, keepdim=True) + var = (xf - mean).pow(2).mean(dim=-1, keepdim=True) + return (xf - mean) * torch.rsqrt(var + EPS) * w.float() + b.float() + + +def rope( + x: torch.Tensor, pos: torch.Tensor, cos_sin: torch.Tensor, interleave: bool +) -> torch.Tensor: + """Apply RoPE to the first ``rot_dim`` elements of x's last dim. + + x: [..., head_dim] fp32. ``cos_sin`` is [max_pos, rot_dim]. ``interleave`` + selects adjacent-pair (GLM) vs split-half NeoX (DeepSeek-V3.2) layout. + """ + rot = cos_sin.shape[-1] + half = rot // 2 + cs = cos_sin[pos.long()] + cos, sin = cs[..., :half], cs[..., half:] + out = x.float().clone() + r = out[..., :rot] + if interleave: + x1, x2 = r[..., 0::2].clone(), r[..., 1::2].clone() + r[..., 0::2] = x1 * cos - x2 * sin + r[..., 1::2] = x2 * cos + x1 * sin + else: + x1, x2 = r[..., :half].clone(), r[..., half:].clone() + r[..., :half] = x1 * cos - x2 * sin + r[..., half:] = x2 * cos + x1 * sin + return out + + +def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize.""" + amax = vals.float().abs().amax(dim=-1, keepdim=True) + scale = torch.clamp(amax, min=1e-4) / FP8_MAX + scale = torch.exp2(torch.ceil(torch.log2(scale))) + q = (vals.float() / scale).to(FP8) + return q, scale.squeeze(-1) + + +def _bf16_ulp(a: torch.Tensor, b: torch.Tensor) -> int: + def key(t): + u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF + return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000) + + return int((key(a) - key(b)).abs().max().item()) + + +def _fp8_ulp(a: torch.Tensor, b: torch.Tensor) -> int: + def key(t): + u = t.contiguous().view(torch.uint8).to(torch.int64) + return torch.where(u >= 0x80, 0xFF - u, u + 0x80) + + return int((key(a) - key(b)).abs().max().item()) + + +def assert_bf16(got: torch.Tensor, ref_fp32: torch.Tensor, msg: str): + # Kernel keeps RMSNorm/RoPE in fp32 and rounds to bf16 once; the fp32 + # reduction/FMA order differs from torch, so a few elements land on the + # opposite side of a round-to-nearest tie. Use the same tolerance the + # sibling deepseek_v4 fused-kernel test uses for this bf16 norm+rope class. + torch.testing.assert_close( + got.float(), ref_fp32.float(), rtol=1e-2, atol=1e-2, msg=lambda m: f"{msg}: {m}" + ) + + +def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str): + assert _fp8_ulp(got, ref) <= 1, f"{msg}: >1 fp8 ULP" + + +# โ”€โ”€ fused_norm_rope โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) +@pytest.mark.parametrize("index_interleave", [True, False]) +@pytest.mark.parametrize("mla_fp8", [False, True]) +def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool): + torch.manual_seed(0) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) + ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # MLA k_pe: interleaved + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + bs = max_pos # single block covering all tokens + mla_dim = KV_LORA + ROPE_DIM + if mla_fp8: + mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8) + mla_dtype = "fp8" + mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32) + else: + mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16) + mla_dtype = "auto" + mla_k_scale = None + idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 # 132 + idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8) + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + ik, + ikw, + ikb, + EPS, + idx_cos_sin, + topk, + slot_mapping=slot, + indexer_k_cache=idx_cache, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype=mla_dtype, + mla_k_scale=mla_k_scale, + has_indexer=True, + index_rope_interleave=index_interleave, + ) + + # q_lora RMSNorm + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm") + + # MLA cache: [kv_c_normed | k_pe_roped(interleaved)] + kv_ref = rms_norm(kv_c, kvw) + kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) + if mla_fp8: + cache = mla_cache.view(FP8)[0, :num_tokens] + s = mla_k_scale.item() + assert_fp8(cache[:, :KV_LORA], (kv_ref / s).to(FP8), "MLA kv fp8") + assert_fp8(cache[:, KV_LORA:], (kpe_ref / s).to(FP8), "MLA k_pe fp8") + else: + cache = mla_cache[0, :num_tokens] + assert_bf16(cache[:, :KV_LORA], kv_ref, "MLA kv bf16") + assert_bf16(cache[:, KV_LORA:], kpe_ref, "MLA k_pe bf16") + + # Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]). + ik_ref = layer_norm(ik, ikw, ikb) + ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave) + q_ref, s_ref = ue8m0_quant(ik_ref) + flat = idx_cache[0].reshape(-1) + vals = flat[: bs * INDEX_HEAD_DIM].view(FP8).reshape(bs, INDEX_HEAD_DIM) + scales = flat[bs * INDEX_HEAD_DIM :].view(torch.float32) + assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8") + torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0) + + # Top-k buffer cleared to -1 on indexer layers. + assert (topk == -1).all(), "topk buffer not cleared on indexer layer" + + +@pytest.mark.parametrize("num_tokens", [1, 17, 512]) +def test_fused_norm_rope_no_indexer(num_tokens: int): + """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched.""" + torch.manual_seed(1) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + bs = max_pos + mla_cache = torch.zeros(1, bs, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16) + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + None, + None, + None, + EPS, + None, + topk, + slot_mapping=slot, + indexer_k_cache=None, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype="auto", + mla_k_scale=None, + has_indexer=False, + index_rope_interleave=False, + ) + + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)") + cache = mla_cache[0, :num_tokens] + assert_bf16(cache[:, :KV_LORA], rms_norm(kv_c, kvw), "MLA kv (no-indexer)") + assert_bf16( + cache[:, KV_LORA:], + rope(k_pe.float(), pos, mla_cos_sin, interleave=True), + "MLA k_pe (no-indexer)", + ) + # Shared layers reuse the previous indexer's top-k: buffer must be untouched. + assert (topk == 7).all(), "topk buffer should be untouched on shared layer" + + +# โ”€โ”€ fused_q โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) +@pytest.mark.parametrize("index_interleave", [True, False]) +def test_fused_q(num_tokens: int, index_interleave: bool): + torch.manual_seed(2) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos + + q_pe = torch.randn( + num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16 + ) + ql_nope = torch.randn( + num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16 + ) + index_q = torch.randn( + num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16 + ) + index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32) + q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32) + softmax_scale = INDEX_HEAD_DIM**-0.5 + head_scale = INDEX_HEADS**-0.5 + q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # q_pe: interleaved + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + iq_fp8, iw_out, mqa = K.fused_q( + pos, + q_pe, + q_cos_sin, + index_q, + idx_cos_sin, + ql_nope, + q_scale, + index_w, + softmax_scale, + head_scale, + has_indexer=True, + index_rope_interleave=index_interleave, + ) + + s = q_scale.item() + # MQA query: [ql_nope | q_pe RoPE'd (interleaved)], per-tensor fp8. + mqa_nope_ref = (ql_nope.float() / s).to(FP8) + qpe_ref = rope( + q_pe.float(), + pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS), + q_cos_sin, + interleave=True, + ) + mqa_pe_ref = (qpe_ref / s).to(FP8) + assert_fp8(mqa[:, :, :KV_LORA], mqa_nope_ref, "mqa ql_nope") + assert_fp8(mqa[:, :, KV_LORA:], mqa_pe_ref, "mqa q_pe") + + # Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale. + iq_ref = rope( + index_q.float(), + pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS), + idx_cos_sin, + interleave=index_interleave, + ) + q_ref, scale_ref = ue8m0_quant(iq_ref) + assert_fp8(iq_fp8, q_ref, "indexer-Q fp8") + iw_ref = index_w * scale_ref * softmax_scale * head_scale + torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("num_tokens", [1, 17, 512]) +def test_fused_q_no_indexer(num_tokens: int): + torch.manual_seed(3) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + q_pe = torch.randn( + num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16 + ) + ql_nope = torch.randn( + num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16 + ) + q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32) + q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + _, _, mqa = K.fused_q( + pos, + q_pe, + q_cos_sin, + None, + None, + ql_nope, + q_scale, + None, + 0.0, + 0.0, + has_indexer=False, + index_rope_interleave=False, + ) + s = q_scale.item() + assert_fp8(mqa[:, :, :KV_LORA], (ql_nope.float() / s).to(FP8), "mqa ql_nope") + qpe_ref = rope( + q_pe.float(), + pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS), + q_cos_sin, + interleave=True, + ) + assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe") + + +# โ”€โ”€ fused_eh_norm (MTP) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) +def test_fused_eh_norm(num_tokens: int): + torch.manual_seed(4) + dev = "cuda" + # Mix in a position-0 token to exercise the embeds-zeroing branch. + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + pos[0] = 0 + embeds = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16) + prev = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16) + ew = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16) + hw = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16) + + out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS) + + masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds) + ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1) + assert out.shape == (num_tokens, 2 * HIDDEN) + assert_bf16(out, ref, "eh_norm") diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b065d2142ce..55c7238f642 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1287,6 +1287,12 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False + # When True, the MoE skips its final cross-rank all-reduce (and the separate + # shared-expert reduce), returning the partial per-rank sum. The caller is + # then responsible for the reduction (e.g. fusing it into the next RMSNorm). + # Only honored on the non-reduced (late-AR) TP path. Default False. + skip_final_all_reduce: bool = False + # SwiGLU clamp limit. When set, backends that do not implement the clamp # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle # cannot silently select one and drop the clamp. 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 b638db13fd2..140466c7f40 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -423,6 +423,7 @@ class MoERunner(MoERunnerInterface): if ( shared_output is not None and not self.moe_config.is_sequence_parallel + and not self.moe_config.skip_final_all_reduce and self._fused_output_is_reduced ): shared_output = tensor_model_parallel_all_reduce(shared_output) @@ -445,6 +446,7 @@ class MoERunner(MoERunnerInterface): # - The MK already reduced the fused output itself. if ( not self.moe_config.is_sequence_parallel + and not self.moe_config.skip_final_all_reduce and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not self._fused_output_is_reduced ): diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index c1bc731ee62..80c0c4ec36a 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -227,6 +227,7 @@ def sparse_attn_indexer( topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, use_fp4_cache: bool = False, + skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -303,7 +304,13 @@ def sparse_attn_indexer( scale_fmt, ) - topk_indices_buffer[: hidden_states.shape[0]] = -1 + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the + # top-k kernels scatter valid indices into it. On the fused deepseek_v32 + # nvidia path, _fused_norm_rope_kernel already cleared the same + # [:num_tokens, :topk] region earlier in this forward, so skip the redundant + # fill. + if not skip_topk_buffer_clear: + topk_indices_buffer[: hidden_states.shape[0]] = -1 if has_prefill: prefill_metadata = attn_metadata_narrowed.prefill assert prefill_metadata is not None @@ -546,6 +553,7 @@ def sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool, use_fp4_cache: bool = False, + skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: return topk_indices_buffer diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index 21b0c2c441d..420e4e93785 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING - import torch import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config @@ -23,7 +21,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_attn_indexer import ( + SparseAttnIndexer, + sparse_attn_indexer, +) from vllm.model_executor.models.deepseek_v2 import ( DeepSeekV2FusedQkvAProjLinear, DeepseekV32IndexerCache, @@ -32,10 +33,7 @@ from vllm.model_executor.models.deepseek_v2 import ( from vllm.model_executor.models.utils import extract_layer_index from vllm.utils.torch_utils import is_quantized_kv_cache -if TYPE_CHECKING: - from vllm.model_executor.layers.attention.mla_attention import ( - MLACommonMetadata, - ) +from .kernels import fused_norm_rope, fused_q class DeepseekV32Indexer(nn.Module): @@ -160,6 +158,10 @@ class DeepseekV32Indexer(nn.Module): class DeepseekV32Attention(MLAAttention): + # Narrow the base's broadly-typed `indexer` to the concrete type so the + # `if self.indexer is not None` guards below type-check its attributes. + indexer: "DeepseekV32Indexer | None" + def __init__( self, vllm_config: VllmConfig, @@ -263,21 +265,27 @@ class DeepseekV32Attention(MLAAttention): self.num_local_heads = num_local_heads self.qk_head_dim = qk_head_dim self.indexer = indexer + self.topk_indices_buffer = topk_indices_buffer # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0 # computes the top-k, steps 1+ set this True to reuse it. self.skip_topk = False - # Whether the paged KV cache must be viewed as fp8 before the attention - # (per-tensor fp8; the fp8_ds_mla layout is read as uint8). - self._fp8_kv_needs_view = ( - is_quantized_kv_cache(self.kv_cache_dtype) - and self.kv_cache_dtype != "fp8_ds_mla" - ) - # Whether the backend takes an fp8-quantized query (FlashInfer sparse) - # vs the (ql_nope, q_pe) tuple (FlashMLA sparse). - self._use_concat_quant = ( + # Single fused fp8 path: Triton fused norm/rope/cache + fused-q write a + # single fp8 MQA query and the contiguous [kv_c; k_pe] MLA cache layout. + # This requires an fp8 KV cache and a sparse MLA backend that accepts a + # quantized query (FlashInfer sparse on SM100). + assert ( is_quantized_kv_cache(self.kv_cache_dtype) and self.impl.supports_quant_query_input + ), ( + "deepseek_v32 (nvidia) requires an fp8 KV cache served by the " + "FlashInfer sparse MLA backend (which accepts a quantized query). " + "Launch with --kv-cache-dtype fp8." ) + # The paged KV cache is stored as uint8 and viewed as fp8 for the decode + # (per-tensor fp8; never the fp8_ds_mla layout on this path). + self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla" + # GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX. + self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) # Remaining MLA projections (registered on this module). self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( @@ -295,10 +303,14 @@ class DeepseekV32Attention(MLAAttention): prefix=f"{prefix}.q_b_proj", ) self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps) + # reduce_results=False: the attention all-reduce is fused with the + # following post_attention_layernorm in the decoder layer via + # fused_allreduce_rms_norm. self.o_proj = RowParallelLinear( num_heads * v_head_dim, hidden_size, bias=False, + reduce_results=False, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -322,102 +334,177 @@ class DeepseekV32Attention(MLAAttention): positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: + # Captured: A-projections (+ indexer A-GEMM on indexer layers). qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] - q_c, kv_lora = qkv_lora.split( - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + q_c, kv_c, k_pe = qkv_lora.split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 ) - q_c = self.q_a_layernorm(q_c) - q = self.q_b_proj(q_c)[0] - kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - kv_c_normed = self.kv_a_layernorm(kv_c) - - q = q.view(-1, self.num_local_heads, self.qk_head_dim) - k_pe = k_pe.unsqueeze(1) - q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( - positions, q[..., self.qk_nope_head_dim :], k_pe - ) + if self.indexer is not None and not self.skip_topk: + kw = self.indexer.wk_weights_proj(hidden_states)[0] + index_k = kw[:, : self.indexer.head_dim] + index_weights = kw[:, self.indexer.head_dim :] + else: + index_k = None + index_weights = None num_tokens = hidden_states.shape[0] - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_nope = q_nope.transpose(0, 1) # (N, B, P) - ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) # (B, N, L) - - # Lightning indexer writes the top-k indices into the shared buffer. - # "Shared" layers (indexer is None) reuse the top-k from the previous - # indexer layer already sitting in the buffer. - if self.indexer is not None and not self.skip_topk: - self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) # type: ignore[operator] - - attn_latent = torch.empty( - (num_tokens, self.num_local_heads, self.kv_lora_rank), - dtype=q.dtype, - device=q.device, - ) - self._sparse_attention(kv_c_normed, k_pe, ql_nope, q_pe, attn_latent) - - # V up-projection + output projection are metadata-independent GEMMs and - # stay captured. output = torch.empty( (num_tokens, self.num_local_heads * self.v_head_dim), - dtype=q.dtype, - device=q.device, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._fused_attention( + positions, q_c, kv_c, k_pe, index_k, index_weights, output ) - self._v_up_proj(attn_latent, out=output) return self.o_proj(output)[0] @eager_break_during_capture - def _sparse_attention( + def _fused_attention( self, - kv_c_normed: torch.Tensor, + positions: torch.Tensor, + q_c: torch.Tensor, + kv_c: torch.Tensor, k_pe: torch.Tensor, - ql_nope: torch.Tensor, - q_pe: torch.Tensor, - attn_latent: torch.Tensor, + index_k: torch.Tensor | None, + index_weights: torch.Tensor | None, + output: torch.Tensor, ) -> None: + # One eager break for the whole attention. In FULL cudagraph mode (pure + # decode) this decorator is a no-op, so everything here is captured; in + # PIECEWISE (prefill) it runs eagerly. The cache writes, sparse indexer, + # and forward_mqa all depend on per-step metadata and must not be split + # out (PIECEWISE capture would otherwise miss them). forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata - attn_metadata: MLACommonMetadata | None if isinstance(attn_metadata_raw, dict): - attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment] + attn_metadata = attn_metadata_raw.get(self.layer_name) elif isinstance(attn_metadata_raw, list): - # Speculative decoding: [0] is the base-model metadata dict. - attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment] + attn_metadata = attn_metadata_raw[0].get(self.layer_name) else: attn_metadata = attn_metadata_raw slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) - self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_c_normed, + mla_slot = slot_mapping.get(self.layer_name) + + if self.indexer is not None: + has_indexer = True + indexer_k_norm_w = self.indexer.k_norm.weight + indexer_k_norm_bias = self.indexer.k_norm.bias + indexer_k_norm_eps = self.indexer.k_norm.variance_epsilon + indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache + indexer_k_cache = self.indexer.k_cache.kv_cache + indexer_softmax_scale = self.indexer.softmax_scale + indexer_n_head_scale = self.indexer.n_head**-0.5 + else: + has_indexer = False + indexer_k_norm_w = None + indexer_k_norm_bias = None + indexer_k_norm_eps = 1e-6 + indexer_k_rope_cos_sin_cache = None + indexer_k_cache = None + indexer_softmax_scale = 0.0 + indexer_n_head_scale = 0.0 + + if attn_metadata is None: + mla_kv_cache = None + mla_k_scale = None + indexer_k_cache = None + mla_slot = None + else: + mla_kv_cache = self.kv_cache + mla_k_scale = self._k_scale + + q_c = fused_norm_rope( + positions, + q_c, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + kv_c, + self.kv_a_layernorm.weight, + self.kv_a_layernorm.variance_epsilon, k_pe, - self.kv_cache, - slot_mapping.get(self.layer_name), - self.kv_cache_dtype, - self._k_scale, + self.rotary_emb.cos_sin_cache, + index_k, + indexer_k_norm_w, + indexer_k_norm_bias, + indexer_k_norm_eps, + indexer_k_rope_cos_sin_cache, + self.topk_indices_buffer, + slot_mapping=mla_slot, + indexer_k_cache=indexer_k_cache, + mla_kv_cache=mla_kv_cache, + mla_kv_cache_dtype=self.kv_cache_dtype, + mla_k_scale=mla_k_scale, + has_indexer=has_indexer, + index_rope_interleave=self._index_rope_interleave, + ) + + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_nope = q_nope.transpose(0, 1) + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) + + if self.indexer is not None: + index_q = self.indexer.wq_b(q_c)[0] + index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim) + else: + index_q = None + + index_q_fp8, index_weights_out, mqa_q = fused_q( + positions, + q_pe, + self.rotary_emb.cos_sin_cache, + index_q, + self.indexer_rope_emb.cos_sin_cache if has_indexer else None, + ql_nope, + self._q_scale, + index_weights, + indexer_softmax_scale, + indexer_n_head_scale, + has_indexer=has_indexer, + index_rope_interleave=self._index_rope_interleave, ) if attn_metadata is None: - # Profile / warmup: zero-fill for DP+EP determinism. - attn_latent.zero_() + output.zero_() return - num_actual = attn_metadata.num_actual_tokens + if self.indexer is not None: + sparse_attn_indexer( + q_c, + self.indexer.k_cache.prefix, + self.indexer.k_cache.kv_cache, + index_q_fp8, + None, # q_scale folded into weights on the fp8 path + None, # k unused when skip_k_cache_insert=True + index_weights_out, + self.indexer.quant_block_size, + self.indexer.scale_fmt, + self.indexer.topk_tokens, + self.indexer.head_dim, + self.indexer.max_model_len, + self.indexer.max_total_seq_len, + self.topk_indices_buffer, + True, # skip_k_cache_insert + False, # use_fp4_cache + True, # skip_topk_buffer_clear (fused_norm_rope already did it) + ) + + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] kv_cache = self.kv_cache if self._fp8_kv_needs_view: kv_cache = kv_cache.view(torch.float8_e4m3fn) - - ql_nope = ql_nope[:num_actual] - q_pe = q_pe[:num_actual] - # FlashInfer sparse takes a single fp8-quantized query; FlashMLA sparse - # takes the (ql_nope, q_pe) tuple and concatenates internally. - mqa_q: torch.Tensor | tuple[torch.Tensor, torch.Tensor] - if self._use_concat_quant: - mqa_q = self._decode_concat_quant_fp8_op(ql_nope, q_pe, self._q_scale) - else: - mqa_q = (ql_nope, q_pe) - - attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] - attn_latent[:num_actual] = attn_out.view( - num_actual, self.num_local_heads, self.kv_lora_rank + attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] + mqa_q[:num_actual], kv_cache, attn_metadata, self ) + x = attn_out.view( + num_actual, self.num_local_heads, self.kv_lora_rank + ).transpose(0, 1) + out = ( + output[:num_actual] + .view(num_actual, self.num_local_heads, self.v_head_dim) + .transpose(0, 1) + ) + torch.bmm(x, self.W_UV, out=out) diff --git a/vllm/models/deepseek_v32/nvidia/fused_ops.py b/vllm/models/deepseek_v32/nvidia/fused_ops.py new file mode 100644 index 00000000000..6a795e2a153 --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/fused_ops.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused ops for deepseek_v32 (eager / breakable-cudagraph path). + +These recover fusions that vLLM's torch.compile passes would normally do but +that don't fire when running eager under the breakable CUDA graph. +""" + +import torch + +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + _AR_RESIDUAL_RMS_NORM, + _can_use_flashinfer, + flashinfer_trtllm_fused_allreduce_norm, +) +from vllm.model_executor.layers.layernorm import RMSNorm + + +def fused_allreduce_rms_norm( + hidden_states: torch.Tensor, + residual: torch.Tensor, + norm: RMSNorm, +) -> tuple[torch.Tensor, torch.Tensor]: + """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer. + + ``hidden_states`` is the per-rank *partial* output of a row-parallel linear + run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right + after. Returns ``(normed_output, new_residual)``, equivalent to + ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit + all-reduce + RMSNorm when the flashinfer fast path is unavailable. + """ + tp_size = get_tensor_model_parallel_world_size() + if tp_size == 1: + return norm(hidden_states, residual) + + if flashinfer_trtllm_fused_allreduce_norm is not None: + ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size) + if ok: + norm_out = torch.empty_like(hidden_states) + # With norm_out provided, the kernel writes the new residual + # (all_reduce(hidden_states) + residual) into the hidden_states + # buffer and the normalized result into norm_out. + flashinfer_trtllm_fused_allreduce_norm( + allreduce_in=hidden_states, + residual=residual, + rms_gamma=norm.weight, + rms_eps=norm.variance_epsilon, + world_size=tp_size, + weight_bias=0.0, # standard RMSNorm (Gemma would use 1.0) + launch_with_pdl=True, + fp32_acc=True, + max_token_num=max_token_num, + pattern_code=_AR_RESIDUAL_RMS_NORM, + norm_out=norm_out, + ) + return norm_out, hidden_states + + reduced = tensor_model_parallel_all_reduce(hidden_states) + return norm(reduced, residual) diff --git a/vllm/models/deepseek_v32/nvidia/kernels.py b/vllm/models/deepseek_v32/nvidia/kernels.py new file mode 100644 index 00000000000..86419b5060d --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/kernels.py @@ -0,0 +1,823 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.triton_utils import tl, triton + +# Cache of tiny 1-element dummy tensors (per device, dtype) reused by the +# has_indexer=False path so the indexer args don't allocate every call. +_DUMMY_CACHE: dict[tuple, torch.Tensor] = {} + + +def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + key = (shape, dtype, device) + t = _DUMMY_CACHE.get(key) + if t is None: + t = torch.empty(shape, dtype=dtype, device=device) + _DUMMY_CACHE[key] = t + return t + + +@triton.jit +def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): + x = x.to(tl.float32) + mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE + rrms = tl.rsqrt(mean_sq + eps) + w = w.to(tl.float32) + return (x * rrms) * w + + +@triton.jit +def _get_cos_sin( + cos_sin_cache_ptr, + cos_sin_cache_stride, + pos, + HALF_ROT_DIM: tl.constexpr, +): + block = tl.arange(0, HALF_ROT_DIM) + cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block) + cos = cos.to(tl.float32) + sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM) + sin = sin.to(tl.float32) + return cos, sin + + +@triton.jit +def _fp8_ue8m0_quantize(vals): + """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale. + + Returns (fp8_vals, scale) so the caller can store them or reuse the scale. + """ + vals = vals.to(tl.float32) + amax = tl.max(tl.abs(vals)) + scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0) + scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale))) + fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv) + return fp8_vals, scale + + +@triton.jit +def _fp8_quant_and_cache_write( + vals, + mask, + slot_idx, + kv_cache_ptr, + kv_cache_scale_ptr, + cache_block_size, + cache_stride, + offsets, + HEAD_DIM: tl.constexpr, +): + k_fp8, scale = _fp8_ue8m0_quantize(vals) + + block_idx = slot_idx // cache_block_size + block_offset = slot_idx % cache_block_size + block_start = block_idx * cache_block_size * cache_stride + + tl.store( + kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets, + k_fp8, + mask=mask, + ) + scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4 + tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale) + + +@triton.jit +def _fused_norm_rope_kernel( + pos_ptr, + # Q RMS norm + q_c_ptr, + q_c_stride, + q_rms_norm_w_ptr, + q_rms_eps, + q_c_out_ptr, + q_c_out_stride, + Q_DIM: tl.constexpr, + Q_BLOCK_SIZE: tl.constexpr, + # KV RMS norm + kv_ptr, + kv_stride, + kv_rms_norm_w_ptr, + kv_rms_eps, + KV_DIM: tl.constexpr, + # KV RoPE + kpe_ptr, + kpe_stride, + kpe_rope_cos_sin_cache_ptr, + kpe_rope_cos_sin_cache_stride, + KPE_HALF_ROT_DIM: tl.constexpr, + # Index K layer norm + index_k_ptr, + index_k_stride, + index_k_layer_norm_w_ptr, + index_k_layer_norm_bias_ptr, + index_k_layer_norm_eps, + INDEX_K_DIM: tl.constexpr, + INDEX_K_BLOCK_SIZE: tl.constexpr, + # Index K RoPE + index_k_rope_cos_sin_cache_ptr, + index_k_rope_cos_sin_cache_stride, + INDEX_K_HALF_ROT_DIM: tl.constexpr, + # Cache params (shared by indexer K and MLA) + slot_mapping_ptr, + # Index K FP8 cache + indexer_cache_ptr, + indexer_cache_scale_ptr, + indexer_cache_block_size, + indexer_cache_stride, + # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr) + mla_cache_ptr, + mla_cache_block_stride, + mla_cache_entry_stride, + MLA_CACHE_FP8: tl.constexpr, + mla_cache_scale_ptr, + # Top k indices + topk_indices_ptr, + topk_indices_stride, + TOPK: tl.constexpr, + TOPK_BLOCK_SIZE: tl.constexpr, + HAS_INDEXER: tl.constexpr, + INDEX_ROPE_INTERLEAVE: tl.constexpr, +): + pid = tl.program_id(0) + tok_idx = tl.program_id(1) + if pid == 3: + if not HAS_INDEXER: + # Shared layer: reuse the previous indexer layer's top-k; do not + # clear the buffer. + return + # Fill top k indices buffer with -1 + for i in range(0, TOPK, TOPK_BLOCK_SIZE): + offset = i + tl.arange(0, TOPK_BLOCK_SIZE) + mask = offset < TOPK + tl.store( + topk_indices_ptr + tok_idx * topk_indices_stride + offset, + -1, + mask=mask, + ) + return + + if slot_mapping_ptr is None: + # Memory profiling run. + return + slot_idx = tl.load(slot_mapping_ptr + tok_idx) + if slot_idx < 0: + # Padding + return + + if pid == 2: + # Q RMS norm + q_block = tl.arange(0, Q_BLOCK_SIZE) + q_mask = q_block < Q_DIM + q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0) + q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask) + q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM) + tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask) + elif pid == 1: + # KV RMS Norm + KV RoPE + MLA concat_and_cache. + # Merged so the normed kv_c and RoPE'd k_pe can be written + # to the MLA KV cache directly without a separate kernel. + + # KV RMS Norm (result stays in registers for MLA cache write) + kv_block = tl.arange(0, KV_DIM) + kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block) + kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block) + kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM) + + # KV RoPE (interleaved) on k_pe โ€” in registers only. + # k_pe is not needed after the cache write (MLA decode reads + # from kv_cache), so we skip writing back to kpe_ptr. + pos = tl.load(pos_ptr + tok_idx) + cos, sin = _get_cos_sin( + kpe_rope_cos_sin_cache_ptr, + kpe_rope_cos_sin_cache_stride, + pos, + KPE_HALF_ROT_DIM, + ) + dim_off = tl.arange(0, KPE_HALF_ROT_DIM) + kpe_base = kpe_ptr + tok_idx * kpe_stride + x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32) + x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32) + r1 = x1 * cos - x2 * sin + r2 = x2 * cos + x1 * sin + + # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache. + if mla_cache_entry_stride == 0: + return + + mla_block_size = mla_cache_block_stride // mla_cache_entry_stride + mla_block_idx = slot_idx // mla_block_size + mla_block_off = slot_idx % mla_block_size + dst = ( + mla_cache_ptr + + mla_block_idx * mla_cache_block_stride + + mla_block_off * mla_cache_entry_stride + ) + # kv_c_normed (KV_DIM elements) + if MLA_CACHE_FP8: + scale = tl.load(mla_cache_scale_ptr) + kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv) + tl.store(dst + kv_block, kv_c_fp8) + else: + tl.store(dst + kv_block, kv_c) + # k_pe_roped (from registers, interleaved layout) + if MLA_CACHE_FP8: + tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv)) + tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv)) + else: + tl.store(dst + KV_DIM + dim_off * 2, r1) + tl.store(dst + KV_DIM + dim_off * 2 + 1, r2) + elif pid == 0: + if not HAS_INDEXER: + # Shared layer: no indexer K to process. + return + # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write. + # Eliminates the separate indexer_k_quant_and_cache kernel launch. + + index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE) + index_k_mask = index_k_block < INDEX_K_DIM + index_k = tl.load( + index_k_ptr + tok_idx * index_k_stride + index_k_block, + mask=index_k_mask, + other=0.0, + ).to(tl.float32) + index_k_w = tl.load( + index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask + ).to(tl.float32) + index_k_b = tl.load( + index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask + ).to(tl.float32) + + # 1. LayerNorm. Keep (mean, rstd) so the RoPE rotation partner can be + # re-normalized in registers below, avoiding a global scratch buffer. + mean = tl.sum(index_k, axis=0) / INDEX_K_DIM + diff = tl.where(index_k_mask, index_k - mean, 0.0) + var = tl.sum(diff * diff, axis=0) / INDEX_K_DIM + rstd = tl.rsqrt(var + index_k_layer_norm_eps) + normed = (index_k - mean) * rstd * index_k_w + index_k_b + + # 2. RoPE on the rotation region. Supports both interleaved (adjacent + # pairs, e.g. GLM-5.2) and NeoX (split-half, e.g. DeepSeek-V3.2). The + # rotation partner is gathered from the read-only inputs and + # re-normalized with the same (mean, rstd) โ€” no scratch, no atomics. + pos = tl.load(pos_ptr + tok_idx) + in_rope = index_k_block < 2 * INDEX_K_HALF_ROT_DIM + if INDEX_ROPE_INTERLEAVE: + # pair i = block // 2; partner = block ^ 1; even -> -sin, odd -> +sin. + cos_idx = index_k_block // 2 + partner_offs = tl.where(in_rope, index_k_block ^ 1, index_k_block) + sign = tl.where(index_k_block % 2 == 0, -1.0, 1.0) + else: + # NeoX: pair across halves; partner = block ^ HALF. + cos_idx = index_k_block % INDEX_K_HALF_ROT_DIM + partner_offs = tl.where( + in_rope, index_k_block ^ INDEX_K_HALF_ROT_DIM, index_k_block + ) + sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0) + cos_full = tl.load( + index_k_rope_cos_sin_cache_ptr + + pos * index_k_rope_cos_sin_cache_stride + + cos_idx, + mask=in_rope, + other=1.0, + ).to(tl.float32) + sin_full = tl.load( + index_k_rope_cos_sin_cache_ptr + + pos * index_k_rope_cos_sin_cache_stride + + INDEX_K_HALF_ROT_DIM + + cos_idx, + mask=in_rope, + other=0.0, + ).to(tl.float32) + # normed[partner_offs] == (raw_partner - mean) * rstd * w_partner + + # b_partner: gather the raw partner and its norm affine (read-only + # loads), then apply the same per-token mean/rstd. + raw_partner = tl.load( + index_k_ptr + tok_idx * index_k_stride + partner_offs, + mask=index_k_mask, + other=0.0, + ).to(tl.float32) + w_partner = tl.load( + index_k_layer_norm_w_ptr + partner_offs, mask=index_k_mask + ).to(tl.float32) + b_partner = tl.load( + index_k_layer_norm_bias_ptr + partner_offs, mask=index_k_mask + ).to(tl.float32) + normed_partner = (raw_partner - mean) * rstd * w_partner + b_partner + roped = normed * cos_full + sign * normed_partner * sin_full + result = tl.where(in_rope, roped, normed) + + # 3. FP8 quantize + cache write from registers. + # No need to write back to index_k_ptr โ€” the only consumer + # (sparse_attn_indexer) reads from the cache, not index_k. + _fp8_quant_and_cache_write( + result, + index_k_mask, + slot_idx, + indexer_cache_ptr, + indexer_cache_scale_ptr, + indexer_cache_block_size, + indexer_cache_stride, + index_k_block, + INDEX_K_DIM, + ) + + +def fused_norm_rope( + positions: torch.Tensor, + q_c: torch.Tensor, + q_rms_norm_w: torch.Tensor, + q_rms_eps: float, + kv_c: torch.Tensor, + kv_rms_norm_w: torch.Tensor, + kv_rms_eps: float, + k_pe: torch.Tensor, + k_rope_cos_sin_cache: torch.Tensor, + index_k: torch.Tensor | None, + index_k_layer_norm_w: torch.Tensor | None, + index_k_layer_norm_bias: torch.Tensor | None, + index_k_layer_norm_eps: float, + index_k_rope_cos_sin_cache: torch.Tensor | None, + topk_indices_buffer: torch.Tensor, + # Cache params for fused writes (single slot_mapping for both caches) + slot_mapping: torch.Tensor | None = None, + indexer_k_cache: torch.Tensor | None = None, + mla_kv_cache: torch.Tensor | None = None, + mla_kv_cache_dtype: str = "auto", + mla_k_scale: torch.Tensor | None = None, + has_indexer: bool = True, + index_rope_interleave: bool = False, + q_c_out: torch.Tensor | None = None, +) -> torch.Tensor: + assert positions.ndim == 1 + assert q_c.ndim == 2 + assert kv_c.ndim == 2 + assert k_pe.ndim == 2 + assert topk_indices_buffer.ndim == 2 + + num_tokens = positions.shape[0] + q_dim = q_c.shape[-1] + kv_dim = kv_c.shape[-1] + device = positions.device + + # Shared (no-indexer) layers: substitute cached 1-element dummies so the + # kernel launches cleanly; pid 0/3 (indexer + topk fill) skipped by + # HAS_INDEXER and never dereference them. + if not has_indexer: + indexer_k_cache = None + index_k = _dummy((1, 1), q_c.dtype, device) + index_k_layer_norm_w = _dummy((1,), torch.float32, device) + index_k_layer_norm_bias = _dummy((1,), torch.float32, device) + index_k_rope_cos_sin_cache = k_rope_cos_sin_cache + assert index_k is not None + assert index_k_rope_cos_sin_cache is not None + index_k_dim = index_k.shape[-1] + topk = topk_indices_buffer.shape[-1] + + # --- Indexer K cache setup --- + if indexer_k_cache is not None: + assert slot_mapping is not None + idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) + idx_cache_block_size = indexer_k_cache.shape[1] + idx_cache_stride = indexer_k_cache.shape[2] + if indexer_k_cache.dtype == torch.uint8: + indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn) + else: + # No indexer cache (shared layer / MLA-only fusion). Use dummies but + # KEEP the caller's slot_mapping so the MLA write (pid 1) still runs. + idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device) + indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device) + idx_cache_block_size = 1 + idx_cache_stride = 1 + if mla_kv_cache is None: + # Pure profiling run (no caches at all): skip all per-token writes. + slot_mapping = torch.full( + (num_tokens,), -1, dtype=torch.int64, device=device + ) + + # --- MLA KV cache setup --- + mla_cache_fp8 = mla_kv_cache_dtype != "auto" + if mla_kv_cache is not None: + mla_block_stride = mla_kv_cache.stride(0) + mla_entry_stride = mla_kv_cache.stride(1) + if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8: + mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn) + if mla_k_scale is None: + mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) + else: + # Dummy values โ€” pid 2 will skip the MLA cache write because + # slot_mapping is all -1. + mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device) + mla_block_stride = 0 + mla_entry_stride = 0 + mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) + + if q_c_out is None: + q_c_out = torch.empty_like(q_c) + _fused_norm_rope_kernel[(4, num_tokens)]( + positions, + # Q RMS norm + q_c, + q_c.stride(0), + q_rms_norm_w, + q_rms_eps, + q_c_out, + q_c_out.stride(0), + q_dim, + triton.next_power_of_2(q_dim), + # KV RMS norm + kv_c, + kv_c.stride(0), + kv_rms_norm_w, + kv_rms_eps, + kv_dim, + # KV RoPE + k_pe, + k_pe.stride(0), + k_rope_cos_sin_cache, + k_rope_cos_sin_cache.stride(0), + k_rope_cos_sin_cache.shape[-1] // 2, + # Index K layer norm + RoPE + FP8 quant + index_k, + index_k.stride(0), + index_k_layer_norm_w, + index_k_layer_norm_bias, + index_k_layer_norm_eps, + index_k_dim, + triton.next_power_of_2(index_k_dim), + index_k_rope_cos_sin_cache, + index_k_rope_cos_sin_cache.stride(0), + index_k_rope_cos_sin_cache.shape[-1] // 2, + # Cache params + slot_mapping, + indexer_k_cache, + idx_cache_scale_view, + idx_cache_block_size, + idx_cache_stride, + # MLA KV cache (uses same slot_mapping) + mla_kv_cache, + mla_block_stride, + mla_entry_stride, + mla_cache_fp8, + mla_k_scale, + # Top k indices buffer + topk_indices_buffer, + topk_indices_buffer.stride(0), + topk, + TOPK_BLOCK_SIZE=1024, + HAS_INDEXER=has_indexer, + INDEX_ROPE_INTERLEAVE=index_rope_interleave, + ) + return q_c_out + + +@triton.jit +def _fused_q_kernel( + pos_ptr, + # MQA query PE: RoPE + FP8 pack into output tail + q_pe_ptr, + q_pe_stride0, + q_pe_stride1, + NUM_Q_HEADS: tl.constexpr, + q_pe_cos_sin_ptr, + q_pe_cos_sin_stride, + Q_PE_HALF_ROT_DIM: tl.constexpr, + # Index Q RoPE + index_q_ptr, + index_q_stride0, + index_q_stride1, + NUM_INDEX_Q_HEADS: tl.constexpr, + index_q_cos_sin_ptr, + index_q_cos_sin_stride, + INDEX_Q_HALF_ROT_DIM: tl.constexpr, + # Index Q Quantize + index_q_fp8_ptr, + index_q_fp8_stride0, + index_q_fp8_stride1, + INDEX_Q_HEAD_DIM: tl.constexpr, + # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8 + ql_nope_ptr, + ql_nope_stride0, + ql_nope_stride1, + mqa_q_fp8_ptr, + mqa_q_fp8_stride0, + mqa_q_fp8_stride1, + q_scale_ptr, + QL_NOPE_DIM: tl.constexpr, + QL_NOPE_BLOCK: tl.constexpr, + # Index weights + index_weights_ptr, + index_weights_stride, + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out_ptr, + index_weights_out_stride, + HAS_INDEXER: tl.constexpr, + INDEX_ROPE_INTERLEAVE: tl.constexpr, +): + pid = tl.program_id(0) + tok_idx = tl.program_id(1) + head_idx = tl.program_id(2) + + if pid == 2: + # ql_nope quantize + pack into the front of mqa_q_fp8. + if 2 * head_idx >= NUM_Q_HEADS: + return + + scale = tl.load(q_scale_ptr) + for local_head in range(2): + q_head_idx = head_idx * 2 + local_head + if q_head_idx < NUM_Q_HEADS: + ql_nope_off = tl.arange(0, QL_NOPE_BLOCK) + ql_nope_mask = ql_nope_off < QL_NOPE_DIM + ql_nope = tl.load( + ql_nope_ptr + + tok_idx * ql_nope_stride0 + + q_head_idx * ql_nope_stride1 + + ql_nope_off, + mask=ql_nope_mask, + ).to(tl.float32) + ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv) + tl.store( + mqa_q_fp8_ptr + + tok_idx * mqa_q_fp8_stride0 + + q_head_idx * mqa_q_fp8_stride1 + + ql_nope_off, + ql_nope_fp8, + mask=ql_nope_mask, + ) + return + elif pid == 0: + # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8. + if 2 * head_idx >= NUM_Q_HEADS: + return + + pos = tl.load(pos_ptr + tok_idx) + cos, sin = _get_cos_sin( + q_pe_cos_sin_ptr, + q_pe_cos_sin_stride, + pos, + Q_PE_HALF_ROT_DIM, + ) + + scale = tl.load(q_scale_ptr) + for local_head in range(2): + q_head_idx = head_idx * 2 + local_head + if q_head_idx < NUM_Q_HEADS: + rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM) + x1 = tl.load( + q_pe_ptr + + tok_idx * q_pe_stride0 + + q_head_idx * q_pe_stride1 + + rot_off * 2, + ).to(tl.float32) + x2 = tl.load( + q_pe_ptr + + tok_idx * q_pe_stride0 + + q_head_idx * q_pe_stride1 + + rot_off * 2 + + 1 + ).to(tl.float32) + r1 = x1 * cos - x2 * sin + r2 = x2 * cos + x1 * sin + tl.store( + mqa_q_fp8_ptr + + tok_idx * mqa_q_fp8_stride0 + + q_head_idx * mqa_q_fp8_stride1 + + QL_NOPE_DIM + + rot_off * 2, + (r1 / scale).to(tl.float8e4nv), + ) + tl.store( + mqa_q_fp8_ptr + + tok_idx * mqa_q_fp8_stride0 + + q_head_idx * mqa_q_fp8_stride1 + + QL_NOPE_DIM + + rot_off * 2 + + 1, + (r2 / scale).to(tl.float8e4nv), + ) + return + elif pid == 1: + # Index Q RoPE + fp8 quant, all in registers. The roped bf16 index_q is + # never consumed (only the fp8 below is), so we avoid an in-place + # store-then-reload round-trip. + if not HAS_INDEXER: + return + if head_idx >= NUM_INDEX_Q_HEADS: + return + + pos = tl.load(pos_ptr + tok_idx) + index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM) + iq_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1 + index_q = tl.load(iq_base + index_q_block).to(tl.float32) + + # RoPE in registers (interleaved for GLM-5.2, NeoX for DeepSeek-V3.2), + # gathering the rotation partner from the read-only input. + in_rope = index_q_block < 2 * INDEX_Q_HALF_ROT_DIM + if INDEX_ROPE_INTERLEAVE: + cos_idx = index_q_block // 2 + partner_offs = tl.where(in_rope, index_q_block ^ 1, index_q_block) + sign = tl.where(index_q_block % 2 == 0, -1.0, 1.0) + else: + cos_idx = index_q_block % INDEX_Q_HALF_ROT_DIM + partner_offs = tl.where( + in_rope, index_q_block ^ INDEX_Q_HALF_ROT_DIM, index_q_block + ) + sign = tl.where(index_q_block < INDEX_Q_HALF_ROT_DIM, -1.0, 1.0) + cos_full = tl.load( + index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + cos_idx, + mask=in_rope, + other=1.0, + ).to(tl.float32) + sin_full = tl.load( + index_q_cos_sin_ptr + + pos * index_q_cos_sin_stride + + INDEX_Q_HALF_ROT_DIM + + cos_idx, + mask=in_rope, + other=0.0, + ).to(tl.float32) + partner = tl.load(iq_base + partner_offs).to(tl.float32) + roped = index_q * cos_full + sign * partner * sin_full + index_q = tl.where(in_rope, roped, index_q) + + # Index Q Quantize (from registers) + index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q) + tl.store( + index_q_fp8_ptr + + tok_idx * index_q_fp8_stride0 + + head_idx * index_q_fp8_stride1 + + index_q_block, + index_q_fp8, + ) + + # Index weights update + index_weights = tl.load( + index_weights_ptr + tok_idx * index_weights_stride + head_idx + ) + index_weights = index_weights.to(tl.float32) + index_weights *= index_q_scale + index_weights *= index_weights_softmax_scale + index_weights *= index_weights_head_scale + tl.store( + index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, + index_weights, + ) + + +def fused_q( + positions: torch.Tensor, + q_pe: torch.Tensor, + q_pe_cos_sin_cache: torch.Tensor, + index_q: torch.Tensor | None, + index_q_cos_sin_cache: torch.Tensor | None, + ql_nope: torch.Tensor, + q_scale: torch.Tensor, + # Index weights + index_weights: torch.Tensor | None, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + has_indexer: bool = True, + index_rope_interleave: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert positions.ndim == 1 + assert q_pe.ndim == 3 + assert q_pe_cos_sin_cache.ndim == 2 + assert ql_nope.ndim == 3 + assert ql_nope.shape[:2] == q_pe.shape[:2] + + num_tokens = positions.shape[0] + num_q_heads = q_pe.shape[1] + # Grid's 3rd dim must cover the MQA-pack heads (pid 0/2 iterate 2 heads + # each) and, when present, the indexer heads (pid 1). + mqa_grid_heads = (num_q_heads + 1) // 2 + if not has_indexer: + # Shared layer: cached 1-element dummies; pid 1 skipped by HAS_INDEXER + # and never dereferences them. + index_q = _dummy((1, 1, 1), q_pe.dtype, q_pe.device) + index_q_cos_sin_cache = q_pe_cos_sin_cache + index_weights = _dummy((1, 1), torch.float32, q_pe.device) + assert index_q is not None and index_q.ndim == 3 + assert index_q_cos_sin_cache is not None + assert index_weights is not None + num_index_q_heads = index_q.shape[1] + index_q_head_dim = index_q.shape[2] + grid_heads = max(mqa_grid_heads, num_index_q_heads) + mqa_q_fp8 = torch.empty( + q_pe.shape[0], + q_pe.shape[1], + ql_nope.shape[2] + q_pe.shape[2], + dtype=torch.float8_e4m3fn, + device=q_pe.device, + ) + + index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn) + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + _fused_q_kernel[(3, num_tokens, grid_heads)]( + positions, + q_pe, + q_pe.stride(0), + q_pe.stride(1), + num_q_heads, + q_pe_cos_sin_cache, + q_pe_cos_sin_cache.stride(0), + q_pe_cos_sin_cache.shape[-1] // 2, + index_q, + index_q.stride(0), + index_q.stride(1), + num_index_q_heads, + index_q_cos_sin_cache, + index_q_cos_sin_cache.stride(0), + index_q_cos_sin_cache.shape[-1] // 2, + index_q_fp8, + index_q_fp8.stride(0), + index_q_fp8.stride(1), + index_q_head_dim, + ql_nope, + ql_nope.stride(0), + ql_nope.stride(1), + mqa_q_fp8, + mqa_q_fp8.stride(0), + mqa_q_fp8.stride(1), + q_scale, + ql_nope.shape[2], + triton.next_power_of_2(ql_nope.shape[2]), + index_weights, + index_weights.stride(0), + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + index_weights_out.stride(0), + HAS_INDEXER=has_indexer, + INDEX_ROPE_INTERLEAVE=index_rope_interleave, + # num_warps=1 is optimal here: each program is a single 128-element + # rope+quant, so the kernel is program-count/occupancy bound, not + # per-program compute bound (swept 1/2/4/8 โ€” 1 wins or ties everywhere). + num_warps=1, + ) + return index_q_fp8, index_weights_out, mqa_q_fp8 + + +@triton.jit +def _fused_eh_norm_kernel( + pos_ptr, + embeds_ptr, + embeds_stride, + prev_ptr, + prev_stride, + enorm_w_ptr, + hnorm_w_ptr, + eps, + out_ptr, + out_stride, + H: tl.constexpr, + BLOCK: tl.constexpr, +): + """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm + and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" + tok = tl.program_id(0) + off = tl.arange(0, BLOCK) + mask = off < H + + pos = tl.load(pos_ptr + tok) + e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0) + e = tl.where(pos == 0, 0.0, e.to(tl.float32)) + ew = tl.load(enorm_w_ptr + off, mask=mask) + e_normed = _rms_norm(e, ew, eps, H) + tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) + + p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0) + hw = tl.load(hnorm_w_ptr + off, mask=mask) + p_normed = _rms_norm(p, hw, eps, H) + tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask) + + +def fused_eh_norm( + positions: torch.Tensor, + inputs_embeds: torch.Tensor, + previous_hidden: torch.Tensor, + enorm_w: torch.Tensor, + hnorm_w: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H].""" + n, h = inputs_embeds.shape + out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) + _fused_eh_norm_kernel[(n,)]( + positions, + inputs_embeds, + inputs_embeds.stride(0), + previous_hidden, + previous_hidden.stride(0), + enorm_w, + hnorm_w, + eps, + out, + out.stride(0), + h, + triton.next_power_of_2(h), + ) + return out diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index dd9e1d65ead..0cd7fb02ed3 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -36,6 +36,7 @@ from vllm.model_executor.models.utils import ( from vllm.sequence import IntermediateTensors from .attention import DeepseekV32Attention +from .fused_ops import fused_allreduce_rms_norm class DeepseekV32DecoderLayer(torch.nn.Module): @@ -77,6 +78,10 @@ class DeepseekV32DecoderLayer(torch.nn.Module): quant_config=quant_config, prefix=f"{prefix}.mlp", ) + # Defer the MoE cross-rank all-reduce; it is fused into the next + # layer's input_layernorm (or the final norm) via + # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner. + self.mlp.experts.moe_config.skip_final_all_reduce = True else: self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, @@ -84,6 +89,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module): hidden_act=config.hidden_act, quant_config=quant_config, prefix=f"{prefix}.mlp", + reduce_results=False, ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( @@ -98,12 +104,23 @@ class DeepseekV32DecoderLayer(torch.nn.Module): residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: if residual is None: + # First layer: hidden_states is the (already reduced) embedding. residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) + # The previous layer's MLP/MoE output is left un-reduced; fuse its + # all-reduce into this input_layernorm. + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) + # self_attn's o_proj runs reduce_results=False; fuse its all-reduce with + # the post-attention RMSNorm. hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's + # input_layernorm (or the model's final norm). hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -201,7 +218,8 @@ class DeepseekV32Model(torch.nn.Module): {"hidden_states": hidden_states, "residual": residual} ) - hidden_states, _ = self.norm(hidden_states, residual) + # Last layer's MoE output is un-reduced; fuse its all-reduce into norm. + hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 482ebecc526..7a04d4f3d37 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - import typing from collections.abc import Callable, Iterable @@ -9,6 +8,7 @@ import torch.nn as nn from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig +from vllm.distributed import tensor_model_parallel_all_reduce from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -35,6 +35,7 @@ from vllm.model_executor.models.utils import ( from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from .kernels import fused_eh_norm from .model import DeepseekV32DecoderLayer @@ -75,15 +76,23 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - previous_hidden_states = self.hnorm(previous_hidden_states) - hidden_states = self.eh_proj( - torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. + eh_input = fused_eh_norm( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, ) + hidden_states = self.eh_proj(eh_input) hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) + # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the + # main model fuses that all-reduce into the next norm, but here the + # recycle hidden is consumed directly, so reduce it now. + hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Return the pre-final-norm recycle hidden (re-fed as the next spec # step's previous_hidden_states); shared_head norm is applied in # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract. From 7544286b04a860fcfb98725345f2b43a5023b844 Mon Sep 17 00:00:00 2001 From: Gonzague de Carpentier <82534773+decarpentierg@users.noreply.github.com> Date: Sun, 28 Jun 2026 17:19:28 +0200 Subject: [PATCH 092/138] [Bugfix] Transformers backend: recompute `mm_token_type_ids` per request for M-RoPE (#46552) Signed-off-by: Gonzague de Carpentier Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../models/transformers/multimodal.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index d111af076da..0f80d569754 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -55,6 +55,8 @@ if TYPE_CHECKING: logger = init_logger(__name__) +_MODALITY_TO_TOKEN_TYPE_ID = {"image": 1, "video": 2, "audio": 3} + class MultiModalProcessingInfo(BaseProcessingInfo): def get_supported_mm_limits(self): @@ -206,9 +208,8 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): ) # For gemma3 we check `token_type_ids` as the key - mm_token_type_ids = processed_data.get( - "mm_token_type_ids", processed_data.pop("token_type_ids", None) - ) + mm_token_type_ids = processed_data.pop("token_type_ids", None) + mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids) # We can infer vLLM style placeholder from token type ids, if we split # it for each input `mm_data`. @@ -377,7 +378,6 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): return None num_image_patches = kwargs.pop("num_image_patches") - kwargs.pop("mm_token_type_ids", None) # used only in `model.get_rope_index` if pixel_values is not None: # ROCm: Force math SDP backend for vision encoder to avoid accuracy issues @@ -468,24 +468,18 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): { "image_grid_thw", "video_grid_thw", - "mm_token_type_ids", "second_per_grid_ts", "audio_feature_lengths", "use_audio_in_video", }, ) - if any( - v - for k, v in kwargs.items() - if k not in {"image_grid_thw", "mm_token_type_ids"} - ): + if any(v for k, v in kwargs.items() if k not in {"image_grid_thw"}): raise NotImplementedError( "Transformers modeling backend only supports images." ) image_grid_thw = kwargs.get("image_grid_thw", []) video_grid_thw = kwargs.get("video_grid_thw", []) - mm_token_type_ids = kwargs.get("mm_token_type_ids") image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)( image_grid_thw @@ -494,8 +488,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): video_grid_thw ) - # In v4 `get_rope_index` doesn't have wildcard `kwargs`, and - # can't accept arbitrary args, even if its value is `None` + # `get_rope_index` doesn't always accept arbitrary `kwargs` kwargs = {} if not hasattr(self, "_get_rope_index_accepts_mm_token_type_ids"): import inspect @@ -507,11 +500,13 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) ) if self._get_rope_index_accepts_mm_token_type_ids: - if mm_token_type_ids: - kwargs["mm_token_type_ids"] = torch.cat(mm_token_type_ids) - else: - shape = (1, len(input_tokens)) - kwargs["mm_token_type_ids"] = torch.zeros(*shape, dtype=torch.int) + mm_token_type_ids = torch.zeros(len(input_tokens), dtype=torch.int) + for feature in mm_features: + position = feature.mm_position + offset, length = position.offset, position.length + mm_token_type_id = _MODALITY_TO_TOKEN_TYPE_ID[feature.modality] + mm_token_type_ids[offset : offset + length] = mm_token_type_id + kwargs["mm_token_type_ids"] = mm_token_type_ids.unsqueeze(0) mrope_positions, mrope_position_delta = self.model.get_rope_index( input_ids=torch.tensor(input_tokens).unsqueeze(0), From 4b643c463e31e0513c4b722c8c0754685159c6fa Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 28 Jun 2026 08:37:00 -0700 Subject: [PATCH 093/138] [GLM5] Fix minor typo (#46961) Signed-off-by: Woosuk Kwon --- vllm/models/deepseek_v32/nvidia/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index 420e4e93785..771a3d4f954 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -392,7 +392,7 @@ class DeepseekV32Attention(MLAAttention): has_indexer = True indexer_k_norm_w = self.indexer.k_norm.weight indexer_k_norm_bias = self.indexer.k_norm.bias - indexer_k_norm_eps = self.indexer.k_norm.variance_epsilon + indexer_k_norm_eps = self.indexer.k_norm.eps indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache indexer_k_cache = self.indexer.k_cache.kv_cache indexer_softmax_scale = self.indexer.softmax_scale From 03c6d01c3028cd567ddddc54e1d8414a24dbe501 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Sun, 28 Jun 2026 19:43:19 +0200 Subject: [PATCH 094/138] [OCP MX ] Add back emulation to available OCP MX backends list (#46629) Signed-off-by: Felix Marty Co-authored-by: Andreas Karatzas --- tests/models/quantization/test_gpt_oss.py | 5 ----- tests/quantization/test_quark.py | 4 ---- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py | 3 +-- 3 files changed, 1 insertion(+), 11 deletions(-) diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py index 1f5e48cb0c2..783f1773d21 100644 --- a/tests/models/quantization/test_gpt_oss.py +++ b/tests/models/quantization/test_gpt_oss.py @@ -104,11 +104,6 @@ def test_gpt_oss_attention_quantization( model_args = EvaluationConfig(model_name).get_model_args(tp_size) - # Emulation backend on MI300, MI250 is opt-in - # following https://github.com/vllm-project/vllm/pull/45896 - if not on_gfx950(): - model_args["moe_backend"] = "emulation" - extra_run_kwargs = { "gen_kwargs": {"max_gen_toks": 8000}, "apply_chat_template": True, diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 38b66552f4e..c1cb18f8e22 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -221,10 +221,6 @@ class AccuracyTestConfig: if model_max_len is not None: model_args["max_model_len"] = model_max_len - # Emulation backend on MI300, MI250 is opt-in following https://github.com/vllm-project/vllm/pull/45896 - if not on_gfx950(): - model_args["moe_backend"] = "emulation" - return model_args diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 56a7a6482d1..b1b41ded11a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -319,6 +319,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, Mxfp4MoeBackend.XPU, + Mxfp4MoeBackend.EMULATION, ] return _AVAILABLE_BACKENDS @@ -554,8 +555,6 @@ def select_mxfp4_moe_backend( f"weight_key=kMxfp4Static, activation_key={activation_key}. " "Native backends require specific hardware. " "Set `VLLM_LOGGING_LEVEL=DEBUG` to see detailed unsupported reasons. " - "To use the emulation backend for research/debugging, pass " - "--moe-backend emulation." ) return Mxfp4MoeBackend.NONE, None From c2127a25c787492fea657b867a6c668a317166fb Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Sun, 28 Jun 2026 12:50:30 -0500 Subject: [PATCH 095/138] [ROCm][CI] Fix `rlhf_async_new_apis` Example On ROCm (#46895) Signed-off-by: Micah Williamson Signed-off-by: Matthew Wong Co-authored-by: Matthew Wong Co-authored-by: Andreas Karatzas --- examples/rl/rlhf_async_new_apis.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index a6adc208860..9c3f4700d9e 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -190,12 +190,11 @@ class TrainModel: # Build platform-specific env vars for Ray -ray_env_vars = { - # Prevent Ray from setting CUDA_VISIBLE_DEVICES - "RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1", -} +ray_env_vars = {} if current_platform.is_rocm(): + # Workaround for RCCL bug. See https://github.com/ROCm/rocm-systems/issues/5756 + ray_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1" # For ROCm, BATCH_INVARIANT vllm is not supported ray_env_vars["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" else: From 95528527eab9077aa4eb1d21ccdc20ef18eb5c95 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sun, 28 Jun 2026 17:36:23 -0400 Subject: [PATCH 096/138] [Bugfix][Mooncake] Fix Mooncake lookup prefixes with DCP > 1 (#46855) Signed-off-by: wzhao18 --- .../unit/test_mooncake_store_worker.py | 48 +++++++++++++++++++ .../kv_connector/v1/mooncake/store/data.py | 6 ++- .../kv_connector/v1/mooncake/store/worker.py | 42 ++++++++++++---- 3 files changed, 86 insertions(+), 10 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index d6ce200f4cf..852c75ae9cc 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1621,6 +1621,54 @@ def _make_bare_worker( return worker +def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): + worker = _make_bare_worker() + worker.tp_size = 4 + worker.num_kv_head = 1 + worker.dcp_size = 4 + worker._init_lookup_key_prefixes() + + assert worker._lookup_expected_per_key == 4 + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0", + "test-model@tp_rank:2@pcp0@dcp2@pp_rank:0@group:0", + "test-model@tp_rank:3@pcp0@dcp3@pp_rank:0@group:0", + ) + + +def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): + worker = _make_bare_worker() + worker.tp_size = 4 + worker.num_kv_head = 1 + worker.pcp_size = 2 + worker.dcp_size = 1 + worker._init_lookup_key_prefixes() + + assert worker._lookup_expected_per_key == 2 + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0", + ) + + +def test_lookup_requires_all_dcp_rank_namespaces(): + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 1 + worker.dcp_size = 4 + worker._init_lookup_key_prefixes() + worker.store.batch_is_exist.return_value = [1, 1, 0, 1] + + assert worker.lookup(16, [b"a0"]) == 0 + assert worker.store.batch_is_exist.call_args.args[0] == [ + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130", + "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0@6130", + "test-model@tp_rank:2@pcp0@dcp2@pp_rank:0@group:0@6130", + "test-model@tp_rank:3@pcp0@dcp3@pp_rank:0@group:0@6130", + ] + + def test_lookup_partial_prefix_returns_first_hit_length(): worker = _make_bare_worker() worker.store.batch_is_exist.return_value = [1, 1, 0] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 7f8168ab382..ef98ec0d4e4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -137,6 +137,8 @@ class PoolKey: key_metadata: KeyMetadata, *, tp_rank: int | None = None, + pcp_rank: int | None = None, + dcp_rank: int | None = None, pp_rank: int | None = None, ) -> str: """Return the stable prefix for a Mooncake pool key.""" @@ -145,8 +147,8 @@ class PoolKey: f"{prefix}" f"{key_metadata.model_name}" f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}" - f"@pcp{key_metadata.pcp_rank}" - f"@dcp{key_metadata.dcp_rank}" + f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}" + f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}" f"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}" f"@group:{key_metadata.group_id}" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index aea2d602e72..e60d2f47a4e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1154,17 +1154,43 @@ class MooncakeStoreWorker: self._init_lookup_key_prefixes() def _init_lookup_key_prefixes(self) -> None: - """Precompute per-group key prefixes expanded across TP/PP ranks.""" - tp_count = min(self.tp_size, self.num_kv_head) + """Prepare per-group key prefixes across parallel rank namespaces.""" + # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces + if self.dcp_size > 1: + # DCP reuses the TP workers and splits each TP group into + # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size. + # Store/load paths do not apply KV-head dedup under DCP + rank_namespaces = tuple( + (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) + for pcp_rank in range(self.pcp_size) + for tp_rank in range(self.tp_size) + for pp_rank in range(self.pp_size) + ) + else: + # Without DCP, TP ranks that share a KV head write identical KV, so + # lookup only needs one TP namespace per unique KV head. + tp_count = min(self.tp_size, self.num_kv_head) + rank_namespaces = tuple( + (tp_rank, pcp_rank, 0, pp_rank) + for pcp_rank in range(self.pcp_size) + for tp_rank in range(tp_count) + for pp_rank in range(self.pp_size) + ) + self._lookup_key_prefixes = tuple( tuple( - PoolKey.build_prefix(db.metadata, tp_rank=tp, pp_rank=pp) - for tp in range(tp_count) - for pp in range(self.pp_size) + PoolKey.build_prefix( + db.metadata, + tp_rank=tp_rank, + pcp_rank=pcp_rank, + dcp_rank=dcp_rank, + pp_rank=pp_rank, + ) + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces ) for db in self.token_dbs ) - self._lookup_expected_per_key = tp_count * self.pp_size + self._lookup_expected_per_key = len(rank_namespaces) def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: """Register a cross-layers KV cache tensor. @@ -1430,12 +1456,12 @@ class MooncakeStoreWorker: def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: """Check how many prefix tokens exist in the store. - Checks across all TP ranks and PP ranks. + Checks across all rank-specific key namespaces that may be loaded. """ if not block_hashes or token_len <= 0: return 0 - # Build per-(group, hash) candidate keys expanded across TP/PP. + # Build per-(group, hash) candidate keys expanded across rank namespaces. # candidate_meta stores the (group, hash_bytes) for key slice. candidate_keys: list[str] = [] candidate_meta: list[tuple[int, bytes]] = [] From 4dfbf1503b4bae722743c483a0079ce2f0633f4c Mon Sep 17 00:00:00 2001 From: Fabian Joswig Date: Mon, 29 Jun 2026 01:18:22 +0200 Subject: [PATCH 097/138] [Model] Add support for openai/privacy-filter (#41026) Signed-off-by: Fabian Joswig Co-authored-by: wang.yuqi Co-authored-by: Tyler Michael Smith --- docs/models/pooling_models/token_classify.md | 1 + .../pooling/test_token_classification.py | 47 +++++++ tests/models/registry.py | 4 + .../layers/fused_moe/oracle/unquantized.py | 9 ++ vllm/model_executor/models/gpt_oss.py | 37 ++++- .../models/openai_privacy_filter.py | 127 ++++++++++++++++++ vllm/model_executor/models/registry.py | 4 + vllm/v1/attention/backends/flash_attn.py | 1 + 8 files changed, 225 insertions(+), 5 deletions(-) create mode 100644 vllm/model_executor/models/openai_privacy_filter.py diff --git a/docs/models/pooling_models/token_classify.md b/docs/models/pooling_models/token_classify.md index 79211846211..6b2cefbde55 100644 --- a/docs/models/pooling_models/token_classify.md +++ b/docs/models/pooling_models/token_classify.md @@ -45,6 +45,7 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m | ------------ | ------ | ----------------- | --------------------------- | --------------------------------------- | | `BertForTokenClassification` | bert-based | `boltuix/NeuroBERT-NER` (see note), etc. | | | | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | | +| `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | | | `Qwen3ForTokenClassification`C | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 412e4721c20..0f993d965c7 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -116,6 +116,53 @@ def test_modernbert_models( torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3) +PRIVACY_FILTER_PROMPTS = [ + "My name is Harry Potter.", + "Email me at harry.potter@hogwarts.edu.", + "Call me on +44 20 7946 0958 tomorrow.", + "My account number is 12345678 and the API key is sk-live-abc123def456.", + "I live at 4 Privet Drive, Little Whinging.", + "Visit https://example.com/profile/harry for more info.", + "We met on 12 January 2024.", +] + + +@pytest.mark.parametrize("model", ["openai/privacy-filter"]) +@pytest.mark.parametrize("dtype", ["bfloat16"]) +@torch.inference_mode +def test_openai_privacy_filter( + hf_runner, + vllm_runner, + model: str, + dtype: str, +) -> None: + with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model: + vllm_outputs = vllm_model.token_classify(PRIVACY_FILTER_PROMPTS) + + hf_model_kwargs = {} + if current_platform.is_rocm(): + hf_model_kwargs["attn_implementation"] = "eager" + + with hf_runner( + model, + dtype=dtype, + auto_cls=AutoModelForTokenClassification, + model_kwargs=hf_model_kwargs, + ) as hf_model: + tokenizer = hf_model.tokenizer + hf_outputs = [] + for prompt in PRIVACY_FILTER_PROMPTS: + inputs = tokenizer([prompt], return_tensors="pt") + inputs = hf_model.wrap_device(inputs) + output = hf_model.model(**inputs) + hf_outputs.append(softmax(output.logits[0])) + + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): + hf_output = hf_output.detach().clone().cpu().float() + vllm_output = vllm_output.detach().clone().cpu().float() + torch.testing.assert_close(hf_output, vllm_output, atol=0.1, rtol=1e-2) + + @pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"]) @pytest.mark.parametrize("dtype", ["float"]) @torch.inference_mode diff --git a/tests/models/registry.py b/tests/models/registry.py index 463ce44851b..0bc68f0f7b2 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -707,6 +707,10 @@ _TOKEN_CLASSIFICATION_EXAMPLE_MODELS = { "ModernBertForTokenClassification": _HfExamplesInfo( "disham993/electrical-ner-ModernBERT-base" ), + "OpenAIPrivacyFilterForTokenClassification": _HfExamplesInfo( + "openai/privacy-filter", + min_transformers_version="5.6.0.dev0", + ), } _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = { diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index a8ed9c1d7fe..6a0dfdb0d60 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -11,6 +11,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops from vllm.config.kernel import MoEBackend from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, ) @@ -83,6 +84,14 @@ def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBac if moe_config.moe_parallel_config.dp_size > 1: _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS) + # HACK: unquantized FlashInfer aliases SWIGLUOAI to plain Swiglu + # (swiglu_alpha/limit only set on the MXFP4 branch). Route to + # Triton's swigluoai_and_mul until that's plumbed through. Same + # demotion pattern as the Qwen3.5/dp_size hack above. + if moe_config.activation == MoEActivation.SWIGLUOAI: + _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM) + _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS) + elif current_platform.is_xpu(): _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU] elif current_platform.is_cpu(): diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 2ff5a9ea79b..01f2752ac54 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -70,6 +70,10 @@ from .utils import ( class OAIAttention(nn.Module): + # Override to switch RoPE convention. gpt-oss uses NeoX (chunk halves); + # privacy-filter and similar derivatives use GPT-J (interleaved pairs). + rope_is_neox_style: bool = True + def __init__( self, config: GptOssConfig, @@ -99,7 +103,7 @@ class OAIAttention(nn.Module): "beta_slow": config.rope_parameters["beta_slow"], "truncate": config.rope_parameters.get("truncate", True), }, - is_neox_style=True, + is_neox_style=self.rope_is_neox_style, ) tp_size = get_tensor_model_parallel_world_size() @@ -133,9 +137,25 @@ class OAIAttention(nn.Module): self.num_local_attention_heads = config.num_attention_heads // tp_size self.num_local_key_value_heads = config.num_key_value_heads // tp_size + self.attn = self._build_attention( + config=config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + def _build_attention( + self, + config: GptOssConfig, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> Attention: + # Override to swap in an encoder-only attention or alter the + # per-layer sliding-window policy. # Only apply sliding window to every other layer sliding_window = config.sliding_window if self.layer_idx % 2 == 0 else None - self.attn = Attention( + return Attention( self.num_local_attention_heads, self.head_dim, self.scaling, @@ -222,6 +242,10 @@ class MLPBlock(torch.nn.Module): class TransformerBlock(torch.nn.Module): + # Override to swap attention/MLP without re-implementing the block. + attention_cls: type[nn.Module] = OAIAttention + mlp_cls: type[nn.Module] = MLPBlock + def __init__( self, vllm_config: VllmConfig, @@ -234,13 +258,13 @@ class TransformerBlock(torch.nn.Module): cache_config = vllm_config.cache_config self.layer_idx = extract_layer_index(prefix) - self.attn = OAIAttention( + self.attn = self.attention_cls( config, prefix=f"{prefix}.attn", quant_config=quant_config, cache_config=cache_config, ) - self.mlp = MLPBlock(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp") + self.mlp = self.mlp_cls(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp") self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5) @@ -266,6 +290,9 @@ class TransformerBlock(torch.nn.Module): @support_torch_compile class GptOssModel(nn.Module, EagleModelMixin): + # Override to swap in an alternative TransformerBlock subclass. + block_cls: type[nn.Module] = TransformerBlock + def __init__( self, *, @@ -282,7 +309,7 @@ class GptOssModel(nn.Module, EagleModelMixin): ) self.start_layer, self.end_layer, self.layers = make_layers( self.config.num_hidden_layers, - lambda prefix: TransformerBlock( + lambda prefix: self.block_cls( vllm_config, prefix=prefix, quant_config=self.quant_config, diff --git a/vllm/model_executor/models/openai_privacy_filter.py b/vllm/model_executor/models/openai_privacy_filter.py new file mode 100644 index 00000000000..4b57544a0fd --- /dev/null +++ b/vllm/model_executor/models/openai_privacy_filter.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only OpenAI Privacy Filter model. + +gpt-oss reused as a bidirectional encoder for token classification: every +layer runs non-causal attention with a banded ยฑsliding_window mask, and +the LM head is replaced with a 33-class BIOES score head. +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.config import CacheConfig, VllmConfig +from vllm.model_executor.layers.attention.encoder_only_attention import ( + EncoderOnlyAttention, +) +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.sequence import IntermediateTensors + +from .gpt_oss import GptOssForCausalLM, GptOssModel, OAIAttention, TransformerBlock +from .interfaces_base import attn_type, default_pooling_type +from .utils import AutoWeightsLoader, maybe_prefix + + +class OpenAIPrivacyFilterAttention(OAIAttention): + # Privacy-filter uses GPT-J style RoPE (interleaved pairs), not NeoX. + rope_is_neox_style = False + + def _build_attention( + self, + config, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> EncoderOnlyAttention: + # HF stores sliding_window+1 so each token attends to ยฑW neighbors; + # the encoder-only path applies this as a symmetric (W-1, W-1) mask. + return EncoderOnlyAttention( + num_heads=self.num_local_attention_heads, + head_size=self.head_dim, + scale=self.scaling, + num_kv_heads=self.num_local_key_value_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=config.sliding_window + 1, + prefix=f"{prefix}.attn", + sinks=self.sinks, + ) + + +class OpenAIPrivacyFilterDecoderLayer(TransformerBlock): + attention_cls = OpenAIPrivacyFilterAttention + + +class OpenAIPrivacyFilterModel(GptOssModel): + block_cls = OpenAIPrivacyFilterDecoderLayer + + +def _interleave_gate_up_concat_to_pairs( + weights: Iterable[tuple[str, torch.Tensor]], +) -> Iterable[tuple[str, torch.Tensor]]: + # HF gate_up_proj is concat [gate | up]; swigluoai_and_mul wants + # gate/up interleaved. MXFP4/quark suffixes are already interleaved. + for name, weight in weights: + if name.endswith(".gate_up_proj") or name.endswith(".gate_up_proj_bias"): + *lead, two_i = weight.shape + i = two_i // 2 + weight = ( + torch.stack([weight[..., :i], weight[..., i:]], dim=-1) + .reshape(*lead, two_i) + .contiguous() + ) + yield name, weight + + +@attn_type("encoder_only") +@default_pooling_type(tok_pooling_type="ALL") +class OpenAIPrivacyFilterForTokenClassification(nn.Module): + is_pooling_model = True + hf_to_vllm_mapper = GptOssForCausalLM.hf_to_vllm_mapper + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.head_dtype = vllm_config.model_config.head_dtype + self.num_labels = config.num_labels + + self.model = OpenAIPrivacyFilterModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.score = nn.Linear( + config.hidden_size, config.num_labels, dtype=self.head_dtype + ) + + pooler_config = vllm_config.model_config.pooler_config + assert pooler_config is not None + self.pooler = pooler_for_token_classify(pooler_config) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + hidden_states = hidden_states.to(self.head_dtype) + return self.score(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights( + _interleave_gate_up_concat_to_pairs(weights), + mapper=self.hf_to_vllm_mapper, + ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index dfc034729d8..efc01033499 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -291,6 +291,10 @@ _TOKEN_CLASSIFICATION_MODELS = { "modernbert", "ModernBertForTokenClassification", ), + "OpenAIPrivacyFilterForTokenClassification": ( + "openai_privacy_filter", + "OpenAIPrivacyFilterForTokenClassification", + ), "Qwen3ASRForcedAlignerForTokenClassification": ( "qwen3_asr_forced_aligner", "Qwen3ASRForcedAlignerForTokenClassification", diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index c167708ac9c..df209794352 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -1164,6 +1164,7 @@ class FlashAttentionImpl(AttentionImpl): k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator] v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator] num_splits=1 if self.batch_invariant_enabled else 0, + s_aux=self.sinks, ) return output From 0472436541c842ecda6d249411f1d35649291a79 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 28 Jun 2026 17:04:01 -0700 Subject: [PATCH 098/138] [Spec Decode] Avoid redundant hidden-states gather in draft prefill (#46968) --- vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 422d3ac6901..f1ab8677f75 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -360,7 +360,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.current_draft_step, self.draft_logits, ) - self.hidden_states[:num_reqs] = hidden_states[last_token_indices] + if last_hidden_states is hidden_states: + self.hidden_states[:num_reqs] = sample_hidden_states + else: + self.hidden_states[:num_reqs] = hidden_states[last_token_indices] self.input_buffers.positions[:num_reqs] = positions def _multi_step_decode( From 311ad689adcde0236d630ca202110f5b0fec85f8 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 29 Jun 2026 01:11:17 +0100 Subject: [PATCH 099/138] Remove boilerplate missed by #46820 (#46956) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/gemma4_unified.py | 2 +- vllm/model_executor/models/mllama4.py | 11 ----------- vllm/model_executor/models/param2moe.py | 20 -------------------- vllm/model_executor/models/sarvam.py | 1 - vllm/models/deepseek_v32/nvidia/model.py | 1 - vllm/models/deepseek_v32/nvidia/mtp.py | 1 - 6 files changed, 1 insertion(+), 35 deletions(-) diff --git a/vllm/model_executor/models/gemma4_unified.py b/vllm/model_executor/models/gemma4_unified.py index 9cc0710c4d0..64c084b6161 100644 --- a/vllm/model_executor/models/gemma4_unified.py +++ b/vllm/model_executor/models/gemma4_unified.py @@ -335,7 +335,6 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration): ) # --- MixtureOfExperts delegation to language_model --- - self.expert_weights = self.language_model.expert_weights self.moe_layers = self.language_model.moe_layers self.num_moe_layers = self.language_model.num_moe_layers self.num_logical_experts = self.language_model.num_logical_experts @@ -345,6 +344,7 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration): self.num_expert_groups = self.language_model.num_expert_groups self.num_shared_experts = self.language_model.num_shared_experts self.num_redundant_experts = self.language_model.num_redundant_experts + self.set_eplb_state = self.language_model.set_eplb_state gen_cfg = vllm_config.model_config.try_get_generation_config() self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 855fe5a47a2..178dae506c6 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -812,17 +812,6 @@ class Llama4ForConditionalGeneration( ) return self.language_model.get_eagle3_default_aux_hidden_state_layers() - def set_eplb_state( - self, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ): - self.language_model.set_eplb_state( - expert_load_view, logical_to_physical_map, logical_replica_count - ) - self.expert_weights = self.language_model.expert_weights - def update_physical_experts_metadata( self, num_physical_experts: int, num_local_physical_experts: int ): diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 3386f9545fa..ff56cf505f0 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -705,8 +705,6 @@ class Param2MoEModel(nn.Module): class Param2MoEMixtureOfExperts(MixtureOfExperts): """Implements the vLLM MixtureOfExperts protocol for Param2MoE.""" - expert_weights: list[torch.Tensor] - def extract_moe_parameters(self, example_moe: Param2MoEMoEBlock | None) -> None: if example_moe is None: raise RuntimeError( @@ -745,24 +743,6 @@ class Param2MoEMixtureOfExperts(MixtureOfExperts): if hasattr(fused, "update_expert_map"): fused.update_expert_map() - def set_eplb_state( - self, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - self.expert_weights = [] - for layer_idx, layer in enumerate(self.moe_layers): - if hasattr(layer, "get_expert_weights"): - self.expert_weights.append(layer.get_expert_weights()) - if hasattr(layer, "set_eplb_state"): - layer.set_eplb_state( - moe_layer_idx=layer_idx, - expert_load_view=expert_load_view, - logical_to_physical_map=logical_to_physical_map, - logical_replica_count=logical_replica_count, - ) - class Param2MoEForCausalLM( nn.Module, SupportsPP, SupportsLoRA, Param2MoEMixtureOfExperts diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index fd28e3b3914..f59579b1bcc 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -711,7 +711,6 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE self.model.make_empty_intermediate_tensors ) - self.expert_weights = [] self.num_moe_layers = 0 self.moe_layers = [] diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 0cd7fb02ed3..b139b8ba23f 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -336,7 +336,6 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): def set_moe_parameters(self): # Same as the base, but keyed on the MoE block type rather than the # decoder-layer type (DeepseekV32DecoderLayer is a plain nn.Module). - self.expert_weights = [] self.num_expert_groups = getattr(self.config, "n_group", 1) self.moe_layers = [] self.moe_mlp_layers = [] diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 7a04d4f3d37..0efa1ac7a7e 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -175,7 +175,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): self.set_moe_parameters() def set_moe_parameters(self): - self.expert_weights = [] self.num_moe_layers = self.config.num_nextn_predict_layers self.num_expert_groups = self.config.n_group self.moe_layers = [] From a2abce646f7db07f2169dfc59433d4128bc404de Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 29 Jun 2026 04:43:58 +0200 Subject: [PATCH 100/138] [EPLB] Mask padding in EPLB load recording (#38128) Signed-off-by: ilmarkov Signed-off-by: Markov Ilya Co-authored-by: Markov Ilya --- .../test_eplb_fused_moe_layer_dep_nvfp4.py | 6 ++ tests/kernels/moe/test_moe_layer.py | 3 + tests/kernels/moe/test_routing.py | 66 +++++++++++++++++ .../test_routed_experts_capture.py | 1 + vllm/config/utils.py | 20 +++++ .../distributed/elastic_ep/elastic_execute.py | 4 +- vllm/distributed/eplb/eplb_state.py | 73 +++++++++++++++++-- .../layers/fused_moe/router/base_router.py | 43 +++++++++-- vllm/models/deepseek_v4/nvidia/model.py | 6 ++ vllm/v1/spec_decode/extract_hidden_states.py | 13 ++++ vllm/v1/spec_decode/llm_base_proposer.py | 18 +++++ vllm/v1/worker/gpu/eplb_utils.py | 12 +++ vllm/v1/worker/gpu/model_runner.py | 3 + .../spec_decode/autoregressive/speculator.py | 4 + .../gpu/spec_decode/dflash/speculator.py | 7 ++ vllm/v1/worker/gpu/spec_decode/speculator.py | 15 ++++ vllm/v1/worker/gpu_model_runner.py | 9 +++ 17 files changed, 289 insertions(+), 14 deletions(-) diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py index 551811e60e8..e2d54821ce9 100644 --- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py +++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py @@ -225,6 +225,12 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): logical_to_physical_map, logical_replica_count, ) + fml.router.eplb_state.should_record_tensor = torch.ones( + (), dtype=torch.bool, device=device + ) + fml.router.eplb_state.num_unpadded_tokens_tensors = [ + torch.tensor(0, dtype=torch.int32, device=device) + ] out_after_shuffle = [] with set_forward_context( diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index d1bcd3241aa..552063988fa 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -1332,6 +1332,9 @@ def _test_body_eplb( eplb_moe_layer.router.eplb_state.should_record_tensor = torch.ones( (), dtype=torch.bool, device=device ) + eplb_moe_layer.router.eplb_state.num_unpadded_tokens_tensors = [ + torch.tensor(0, dtype=torch.int32, device=device) + ] # Get "after" output with rearranged weights and EPLB routing with set_forward_context( diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 41dea812193..62a4968a0d1 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -61,12 +61,14 @@ def setup_eplb_state( global_num_experts, dtype=torch.int64, device="cuda" ) should_record_tensor = torch.ones((), dtype=torch.bool, device="cuda") + num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32, device="cuda")] return EplbLayerState( expert_load_view=expert_load_view, logical_to_physical_map=logical_to_physical_map, logical_replica_count=logical_replica_count, should_record_tensor=should_record_tensor, + num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) @@ -782,3 +784,67 @@ def test_eplb_map_with_redundancy( torch.testing.assert_close(load, exp_load) else: assert load.sum().item() == 0 + + +@pytest.mark.parametrize( + "l2p_map, replica_count, num_physical, topk_ids, " + "num_unpadded, expected_out, expected_load", + [ + pytest.param( + [[0], [1], [2], [3]], + [1, 1, 1, 1], + 4, + [[0, 1], [2, 3], [0, 2], [1, 3]], + 2, + [[0, 1], [2, 3], [0, 2], [1, 3]], + # only rows 0,1 counted: expert 0โ†’1, 1โ†’1, 2โ†’1, 3โ†’1 + [1, 1, 1, 1], + id="half_padded", + ), + pytest.param( + # record everything (None = no padding info) + [[0], [1], [2], [3]], + [1, 1, 1, 1], + 4, + [[0, 1], [2, 3], [0, 2], [1, 3]], + None, + [[0, 1], [2, 3], [0, 2], [1, 3]], + [2, 2, 2, 2], + id="no_padding_info", + ), + ], +) +def test_eplb_map_num_unpadded_tokens( + l2p_map, + replica_count, + num_physical, + topk_ids, + num_unpadded, + expected_out, + expected_load, +): + l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda") + rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda") + load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") + rec = torch.tensor(True, dtype=torch.bool, device="cuda") + ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda") + num_unpadded_t = ( + torch.tensor(num_unpadded, dtype=torch.int32, device="cuda") + if num_unpadded is not None + else None + ) + + out = eplb_map_to_physical_and_record( + topk_ids=ids, + expert_load_view=load, + logical_to_physical_map=l2p, + logical_replica_count=rc, + record_enabled=rec, + num_unpadded_tokens=num_unpadded_t, + ) + + exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda") + torch.testing.assert_close(out, exp_out) + + exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda") + torch.testing.assert_close(load, exp_load) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index d1a542396e6..9efee9eec82 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -91,6 +91,7 @@ def test_base_router_capture_with_eplb_enabled(): eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1) eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64) eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool) + eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)] router = _make_router(eplb_state=eplb_state) captured = [] diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 12e0385aeb1..3df0f7210f7 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -203,6 +203,26 @@ class SupportsHash(Protocol): def compute_hash(self) -> str: ... +_config_hash_cache: dict[int, str] = {} + + +def compute_hash_cached(config: SupportsHash) -> str: + """Cache config.compute_hash() by object identity. + + Config objects (ModelConfig, etc.) are long-lived singletons that never + mutate after construction, but compute_hash() is expensive (JSON + serialization + SHA-256). This utility avoids recomputing the hash on + every forward pass while keeping a single consistent key type for all + lookup paths. + """ + key = id(config) + result = _config_hash_cache.get(key) + if result is None: + result = config.compute_hash() + _config_hash_cache[key] = result + return result + + class SupportsMetricsInfo(Protocol): def metrics_info(self) -> dict[str, str]: ... diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 3cb0d603e3e..b0c3740f57e 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -458,7 +458,9 @@ class ElasticEPScalingExecutor: eplb_model_state.logical_to_physical_map, eplb_model_state.logical_replica_count, ) - eplb_state._init_should_record_tensor(model) + eplb_state._propagate_shared_tensors( + model, eplb_model_state.num_unpadded_tokens_tensors + ) model.update_physical_experts_metadata( num_physical_experts=num_physical_experts, num_local_physical_experts=num_local_experts, diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 74f357fbdbf..feacb03d28b 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -35,6 +35,7 @@ import torch from torch.distributed import ProcessGroup, all_reduce from vllm.config import ModelConfig, ParallelConfig +from vllm.config.utils import compute_hash_cached from vllm.distributed.parallel_state import ( get_ep_group, get_eplb_group, @@ -206,6 +207,13 @@ class EplbModelState: pending_result relies on the GIL to synchronize access between the main thread and the async worker. """ + num_unpadded_tokens_tensors: list[torch.Tensor] | None = None + """ + Per-ubatch scalar int32 tensors holding the number of real (non-padding) + tokens. Allocated once in :meth:`EplbState.add_model` so that device + pointers remain stable across CUDA-graph replays. The router kernel + indexes this list with ``dbo_current_ubatch_id()``. + """ class EplbState: @@ -253,7 +261,7 @@ class EplbState: Shared scalar bool tensor for all layers. Every :class:`EplbLayerState` holds a reference to the **same** object so a single ``.fill_()`` updates all layers at once. Allocated on the - first call to :meth:`_init_should_record_tensor`. + first call to :meth:`_propagate_shared_tensors`. """ self.is_async: bool = False """ @@ -440,12 +448,19 @@ class EplbState: self.policy = EPLB_POLICIES[policy_type] logger.debug("Selected EPLB policy: %s", policy_type) + # num_ubatches is 0 when DBO is disabled. + num_ubatches = max(1, self.parallel_config.num_ubatches) + num_unpadded_tokens_tensors = [ + torch.tensor(0, dtype=torch.int32, device=self.device) + for _ in range(num_ubatches) + ] + model.set_eplb_state( expert_load_pass, logical_to_physical_map, logical_replica_count, ) - self._init_should_record_tensor(model) + self._propagate_shared_tensors(model, num_unpadded_tokens_tensors) expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]] assert self.parallel_config.eplb_config.communicator is not None, ( @@ -471,10 +486,43 @@ class EplbState: eplb_stats=None, cuda_device_index=self.cuda_device_index, communicator=communicator, + num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) self.model_states[model_config.compute_hash()] = model_state self.num_valid_physical_experts = model.num_physical_experts + def prepare_forward( + self, + model_config: ModelConfig, + num_unpadded_tokens: int, + ubatch_slices: list | None = None, + ) -> None: + """Fill the per-[u]batch ``num_unpadded_tokens`` tensors before a + forward pass. + + Args: + model_config: Identifies which ``EplbModelState`` to update. + num_unpadded_tokens: Total number of real (non-padding) tokens + in the batch. + ubatch_slices: When DBO is active, a list of + ``UBatchSlice`` objects describing each micro-batch's + token range. When ``None``, only ``tensors[0]`` is filled. + """ + model_state = self.model_states.get(compute_hash_cached(model_config)) + if model_state is None or model_state.num_unpadded_tokens_tensors is None: + return + tensors = model_state.num_unpadded_tokens_tensors + if ubatch_slices is None: + tensors[0].fill_(num_unpadded_tokens) + else: + for i, ubatch_slice in enumerate(ubatch_slices): + ts = ubatch_slice.token_slice + # Real tokens in this ubatch: clamp the global count into + # the slice range so partially-filled ubatches get the + # correct count. + val = max(0, min(num_unpadded_tokens, ts.stop) - ts.start) + tensors[i].fill_(val) + def step( self, is_dummy: bool = False, @@ -638,11 +686,20 @@ class EplbState: self._should_record_current_step(log_stats=log_stats) ) - def _init_should_record_tensor(self, model: "MixtureOfExperts") -> None: # type: ignore[name-defined] - """Allocate (once) and propagate the shared ``should_record_tensor``. + def _propagate_shared_tensors( + self, + model: "MixtureOfExperts", # type: ignore[name-defined] + num_unpadded_tokens_tensors: list[torch.Tensor], + ) -> None: + """Propagate shared tensors to every :class:`EplbLayerState`. + + Allocates ``should_record_tensor`` on the first call and then + assigns both it and ``num_unpadded_tokens_tensors`` to every + MoE layer's :class:`EplbLayerState`. All layers reference the + **same** objects so a single update is visible everywhere. Must be called after :meth:`model.set_eplb_state` so that each - layer's ``eplb_state`` is already populated with the tensor views. + layer's ``eplb_state`` is already populated. """ layer_states = [ layer.eplb_state @@ -659,6 +716,7 @@ class EplbState: for ls in layer_states: if ls is not None: ls.should_record_tensor = self.should_record_tensor + ls.num_unpadded_tokens_tensors = num_unpadded_tokens_tensors def rearrange( self, @@ -985,6 +1043,11 @@ class EplbLayerState: sliding window before the next rearrangement, so recording them wastes GPU work. """ + num_unpadded_tokens_tensors: list[torch.Tensor] | None = None + """ + Reference to the parent :class:`EplbModelState`'s tensor list so the + router can read the correct per-[u]batch unpadded token count. + """ def set_layer_state( self, diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index 4ba855b645f..01e674b2b13 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.v1.worker.ubatching import dbo_current_ubatch_id if current_platform.is_cuda_alike(): @@ -22,11 +23,13 @@ if current_platform.is_cuda_alike(): out_ids_ptr, out_ptr, record_enabled_ptr, + num_unpadded_tokens_ptr, num_logical_experts, map_slots, out_size, numel, num_active_experts, + HAS_NUM_UNPADDED: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) @@ -50,6 +53,13 @@ if current_platform.is_cuda_alike(): token_idx = (offs // num_active_experts).to(tl.int64) hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF replica_idx = hashed % replica_count + map_index = safe_expert_id * map_slots + replica_idx + physical_id = tl.load( + logical_to_physical_ptr + map_index, + mask=mask & valid_expert, + other=-1, + ) + tl.store(out_ids_ptr + offs, physical_id, mask=mask) # 2. Record expert load metrics. @@ -64,16 +74,21 @@ if current_platform.is_cuda_alike(): # If later refactor moved all the MoE kernel calls # to the modular kernel, we can move this logic there # to achieve better efficiency. - map_index = safe_expert_id * map_slots + replica_idx - physical_id = tl.load( - logical_to_physical_ptr + map_index, - mask=mask & valid_expert, - other=-1, - ) - tl.store(out_ids_ptr + offs, physical_id, mask=mask) record_enabled = tl.load(record_enabled_ptr) != 0 - valid = mask & record_enabled & (physical_id >= 0) & (physical_id < out_size) + # Skip padded tokens when recording. + if HAS_NUM_UNPADDED: + num_unpadded_tokens = tl.load(num_unpadded_tokens_ptr) + is_unpadded = offs < num_unpadded_tokens * num_active_experts + else: + is_unpadded = True + valid = ( + mask + & record_enabled + & is_unpadded + & (physical_id >= 0) + & (physical_id < out_size) + ) safe_physical_id = tl.where(physical_id >= 0, physical_id, 0) tl.atomic_add(out_ptr + safe_physical_id, 1, mask=valid) @@ -83,6 +98,7 @@ if current_platform.is_cuda_alike(): logical_replica_count: torch.Tensor, expert_load_view: torch.Tensor, record_enabled: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None, ) -> torch.Tensor: topk_ids_in = topk_ids.contiguous().to(dtype=torch.int32) numel = topk_ids_in.numel() @@ -99,11 +115,13 @@ if current_platform.is_cuda_alike(): out_flat, expert_load_view, record_enabled, + num_unpadded_tokens, logical_replica_count.shape[0], logical_to_physical_map.shape[1], expert_load_view.shape[0], numel, num_active_experts, + HAS_NUM_UNPADDED=num_unpadded_tokens is not None, BLOCK_SIZE=256, ) return out_flat.reshape(topk_ids.shape) @@ -114,6 +132,7 @@ if current_platform.is_cuda_alike(): logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, record_enabled: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None = None, ) -> torch.Tensor: # Fused triton implementation: mapping + optional recording in one kernel. return _eplb_map_and_record_triton( @@ -122,6 +141,7 @@ if current_platform.is_cuda_alike(): logical_replica_count=logical_replica_count, expert_load_view=expert_load_view, record_enabled=record_enabled, + num_unpadded_tokens=num_unpadded_tokens, ) else: @@ -131,6 +151,7 @@ else: logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, record_enabled: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None = None, ) -> torch.Tensor: return topk_ids @@ -177,6 +198,8 @@ class BaseRouter(FusedMoERouter): raise ValueError("EPLB requires logical_replica_count != None") if eplb_state.should_record_tensor is None: raise ValueError("EPLB requires should_record_tensor != None") + if eplb_state.num_unpadded_tokens_tensors is None: + raise ValueError("EPLB requires num_unpadded_tokens_tensors != None") def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor: """Apply EPLB mapping to convert logical expert IDs to physical expert IDs.""" @@ -186,12 +209,16 @@ class BaseRouter(FusedMoERouter): assert eplb_state.logical_to_physical_map is not None assert eplb_state.logical_replica_count is not None assert eplb_state.should_record_tensor is not None + assert eplb_state.num_unpadded_tokens_tensors is not None return eplb_map_to_physical_and_record( topk_ids=topk_ids, logical_to_physical_map=eplb_state.logical_to_physical_map, logical_replica_count=eplb_state.logical_replica_count, expert_load_view=eplb_state.expert_load_view, record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ], ) return topk_ids diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 99373361922..f1bcd534e97 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -70,6 +70,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.worker.ubatching import dbo_current_ubatch_id class DeepseekV4MLP(nn.Module): @@ -464,6 +465,11 @@ class DeepseekV4MegaMoEExperts(nn.Module): logical_to_physical_map=eplb_state.logical_to_physical_map, logical_replica_count=eplb_state.logical_replica_count, record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, ) prepare_megamoe_inputs( diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index b6f9eac4dfa..de7a075e2f7 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -9,6 +9,7 @@ import torch import torch.nn as nn from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config +from vllm.distributed.eplb.eplb_state import EplbState from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.model_loader import get_model @@ -43,6 +44,8 @@ class ExtractHiddenStatesProposer: self.dtype = vllm_config.model_config.dtype self.dp_rank = vllm_config.parallel_config.data_parallel_rank + self.eplb_state: EplbState | None = None + # Model and attention layer tracking (initialized in load_model) self.model: nn.Module | None = None self.attn_layer_names: list[str] = [] @@ -83,6 +86,10 @@ class ExtractHiddenStatesProposer: self.max_num_tokens, dtype=torch.int64, device=device ) + def set_eplb_state(self, eplb_state: EplbState) -> None: + """Inject EPLB state after construction.""" + self.eplb_state = eplb_state + def propose( self, num_speculative_tokens: int, @@ -145,6 +152,12 @@ class ExtractHiddenStatesProposer: if num_tokens_across_dp is not None: num_tokens_across_dp[self.dp_rank] = num_input_tokens + if self.eplb_state is not None: + assert self.vllm_config.speculative_config is not None + self.eplb_state.prepare_forward( + self.vllm_config.speculative_config.draft_model_config, + num_tokens, + ) with set_forward_context( per_layer_attn_metadata, self.vllm_config, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index c78d0660665..4eaf6e9e4f8 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -15,6 +15,7 @@ from vllm.config import ( get_layers_from_vllm_config, replace, ) +from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger @@ -79,6 +80,7 @@ class SpecDecodeBaseProposer: self.dtype = vllm_config.model_config.dtype self.max_model_len = vllm_config.model_config.max_model_len self.dp_rank = vllm_config.parallel_config.data_parallel_rank + self.eplb_state: EplbState | None = None self.num_speculative_tokens = self.speculative_config.num_speculative_tokens # We need to get the hidden size from the draft model config because @@ -328,6 +330,10 @@ class SpecDecodeBaseProposer: "does not support M-RoPE yet" ) + def set_eplb_state(self, eplb_state: EplbState) -> None: + """Inject EPLB state after construction.""" + self.eplb_state = eplb_state + def _init_parallel_drafting_params(self): # For parallel drafting, we need the token ID to use for masked slots # And for EAGLE + parallel drafting, we need the hidden state tensor to use @@ -527,6 +533,12 @@ class SpecDecodeBaseProposer: if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"): self.model.model.set_skip_topk(False) + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.draft_model_config, + num_tokens, + ) + with set_forward_context( per_layer_attn_metadata, self.vllm_config, @@ -672,6 +684,12 @@ class SpecDecodeBaseProposer: if self.pass_hidden_states_to_model: model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size] + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.draft_model_config, + batch_size, + ) + with set_forward_context( per_layer_attn_metadata, self.vllm_config, diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index 8f04ce3577c..aea6fdeff83 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -8,6 +8,7 @@ from typing import Any import torch import torch.nn as nn +from vllm.config import ModelConfig from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( @@ -90,6 +91,7 @@ class EPLBController: draft_model, speculative_config.draft_model_config, ) + speculator.set_eplb_state(self.state) self._has_registered_models = True return True @@ -135,6 +137,16 @@ class EPLBController: log_stats=self.parallel_config.eplb_config.log_balancedness, ) + def prepare_forward( + self, + model_config: ModelConfig, + num_unpadded_tokens: int, + ubatch_slices: list | None = None, + ) -> None: + if self.state is None or not self.parallel_config.enable_eplb: + return + self.state.prepare_forward(model_config, num_unpadded_tokens, ubatch_slices) + def setup_from_mapping( self, model: nn.Module, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ce1bb7f5504..0f57e8a31cd 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1284,6 +1284,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) del intermediate_tensors + # Update the EPLB meta. + self.eplb.prepare_forward(self.model_config, input_batch.num_tokens) + # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index f1ab8677f75..747fb3a3905 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -213,6 +213,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): need_eager=is_profile, ) + self._prepare_eplb_forward(input_batch.num_tokens) + if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. assert self.prefill_cudagraph_manager is not None @@ -424,6 +426,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, ) -> None: + self._prepare_eplb_forward(num_reqs) + idx_mapping = self.idx_mapping[:num_reqs] positions = self.input_buffers.positions[:num_reqs] # Run the draft model forward pass. diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 1bd130838a1..e4583967492 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -278,6 +278,9 @@ class DFlashSpeculator(DraftModelSpeculator): self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], ) + # DFlash processes all speculative tokens in one forward pass, + # so the real token count is num_query_tokens. + self._prepare_eplb_forward(num_query_tokens) self._generate_draft( num_reqs, num_query_tokens, @@ -354,6 +357,10 @@ class DFlashSpeculator(DraftModelSpeculator): self.kv_cache_config, ) + # DFlash processes all speculative tokens in one forward pass, + # so the real token count is num_query_tokens. + self._prepare_eplb_forward(num_query_tokens) + if batch_desc.cg_mode == CUDAGraphMode.FULL: assert self.query_cudagraph_manager is not None self.query_cudagraph_manager.run_fullgraph(batch_desc) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index b06c9372a95..341ed715c7a 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -8,6 +8,7 @@ import torch.nn as nn from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.v1.kv_cache_interface import KVCacheConfig @@ -106,6 +107,8 @@ class DraftModelSpeculator(BaseSpeculator): self.dp_size = vllm_config.parallel_config.data_parallel_size self.dp_rank = vllm_config.parallel_config.data_parallel_rank + self.eplb_state: EplbState | None = None + self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -165,6 +168,18 @@ class DraftModelSpeculator(BaseSpeculator): ) self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names + def set_eplb_state(self, eplb_state: EplbState) -> None: + """Inject EPLB state after construction.""" + self.eplb_state = eplb_state + + def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: + """Call EPLB prepare_forward if EPLB is active for the draft model.""" + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.speculative_config.draft_model_config, + num_unpadded_tokens, + ) + def set_attn( self, model_state: ModelState, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 6af53115775..ff1eba09fd0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4319,6 +4319,13 @@ class GPUModelRunner( # When spec decode is enabled, defer connector finalization # (wait_for_save + clear metadata) until after draft model runs. defer_kv_connector_finalize = self.speculative_config is not None + # Update the EPLB meta. + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.model_config, + num_tokens_unpadded, + ubatch_slices_padded, + ) with ( set_forward_context( attn_metadata, @@ -5215,6 +5222,8 @@ class GPUModelRunner( self.drafter.model, spec_config.draft_model_config, ) + assert hasattr(self.drafter, "set_eplb_state") + self.drafter.set_eplb_state(self.eplb_state) eplb_models += 1 self._setup_eagle3_aux_hidden_state_outputs() From 58d6a6e60ae6bd94a20ea6da27eb224188b24dca Mon Sep 17 00:00:00 2001 From: Yuwen Zhou Date: Mon, 29 Jun 2026 11:04:05 +0800 Subject: [PATCH 101/138] [CPU] Support cpu compressed-tensor w8a8 int8 moe (#42920) Signed-off-by: yuwenzho Signed-off-by: Yuwen Zhou --- .buildkite/hardware_tests/cpu.yaml | 4 +- tests/kernels/moe/test_cpu_quant_fused_moe.py | 143 +++++++++++++- tests/quantization/test_cpu_w8a8.py | 22 +++ .../layers/fused_moe/experts/cpu_moe.py | 174 +++++++++++++++++- .../layers/fused_moe/oracle/int8.py | 40 +++- .../compressed_tensors_moe_w8a8_int8.py | 36 +++- 6 files changed, 410 insertions(+), 9 deletions(-) create mode 100644 tests/quantization/test_cpu_w8a8.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index dd85400f2f1..6f1bd344540 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -91,11 +91,13 @@ steps: - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py - tests/quantization/test_compressed_tensors.py - tests/quantization/test_cpu_wna16.py + - tests/quantization/test_cpu_w8a8.py commands: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m " pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs - pytest -x -v -s tests/quantization/test_cpu_wna16.py" + pytest -x -v -s tests/quantization/test_cpu_wna16.py + pytest -x -v -s tests/quantization/test_cpu_w8a8.py" - label: CPU-Distributed Tests (PP+TP) depends_on: [] diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index d8c1b9f2cb6..e0e0203c6b2 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for CPU quantized fused MoE kernels (FP8 W8A16 and MXFP4 W4A16).""" +"""Tests for CPU quantized fused MoE kernels.""" import math import sys @@ -31,7 +31,10 @@ def _prepack_experts(w: torch.Tensor) -> torch.Tensor: return torch.ops._C.convert_weight_packed(w) -# FP8 W8A16 block-scaled fused MoE +# =========================================================================== +# FP8 W8A16 MoE +# =========================================================================== + BLOCK_SIZE = [128, 128] # [block_n, block_k] @@ -216,7 +219,9 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed): torch.testing.assert_close(out_inplace, out, atol=0, rtol=0) -# MXFP4 W4A16 fused MoE +# =========================================================================== +# MXFP4 W4A16 MoE +# =========================================================================== class MXFP4QuantizeUtil: @@ -496,7 +501,9 @@ def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed): torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) -# INT4 W4A16 group-quantized MoE +# =========================================================================== +# INT4 W4A16 MoE +# =========================================================================== def _pack_int4_gptq(w_int4: torch.Tensor) -> torch.Tensor: @@ -749,5 +756,133 @@ def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) +# =========================================================================== +# INT8 W8A8 MoE +# =========================================================================== + + +def _quantize_per_channel(w): + """Symmetric per-channel INT8 quantisation. w: [N, K] -> (int8, scale).""" + amax = w.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + scale = amax / 127.0 + w_q = (w / scale).round().clamp(-128, 127).to(torch.int8) + return w_q, scale.float() + + +def _quantize_per_token(x): + """Symmetric per-token INT8 quantisation. x: [M, K] -> (int8, scale).""" + amax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + scale = amax / 127.0 + x_q = (x / scale).round().clamp(-128, 127).to(torch.int8) + return x_q, scale.float() + + +def _ref_int8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids): + """Reference INT8 W8A8 per-channel fused MoE in pure torch.""" + B, D = a.shape + topk = topk_ids.size(1) + + out = torch.zeros(B, topk, w2.shape[1], dtype=torch.float32) + for b in range(B): + for t in range(topk): + eid = topk_ids[b, t].item() + + x = a[b : b + 1].float() + x_q, x_s = _quantize_per_token(x) + ic = torch.matmul(x_q.float(), w1[eid].float().t()) + ic = ic * x_s * w1_s[eid].view(1, -1) + ic = _silu_and_mul(ic) + + ic_q, ic_s = _quantize_per_token(ic) + oc = torch.matmul(ic_q.float(), w2[eid].float().t()) + oc = oc * ic_s * w2_s[eid].view(1, -1) + out[b, t] = oc.squeeze(0) + + result = (out * topk_weight.unsqueeze(-1)).sum(dim=1) + return result.to(a.dtype) + + +def _make_int8_moe_weights(E, N, K): + factor = 1e-2 + w1_f = (torch.randn(E, 2 * N, K) - 0.5) * 2 + w2_f = (torch.randn(E, K, N) - 0.5) * 2 + + w1_q_list, w1_s_list = [], [] + w2_q_list, w2_s_list = [], [] + for e in range(E): + q, s = _quantize_per_channel(w1_f[e]) + w1_q_list.append(q) + w1_s_list.append(s) + q, s = _quantize_per_channel(w2_f[e]) + w2_q_list.append(q) + w2_s_list.append(s) + + return ( + torch.stack(w1_q_list), + torch.stack(w2_q_list), + torch.stack(w1_s_list) * factor, + torch.stack(w2_s_list) * factor, + ) + + +INT8_NUM_TOKENS = [1, 2, 64, 121] +INT8_MOE_CONFIGS = [ + # (N, K, E, topk) + (256, 512, 8, 2), + (512, 256, 8, 2), + (512, 512, 8, 4), + (768, 2048, 8, 2), +] + + +@pytest.mark.parametrize("M", INT8_NUM_TOKENS) +@pytest.mark.parametrize("N,K,E,topk", INT8_MOE_CONFIGS) +@pytest.mark.parametrize("seed", [0]) +@pytest.mark.parametrize("is_vnni", [False, True]) +@pytest.mark.parametrize("inplace", [False, True]) +def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace): + """Test fused_experts_cpu INT8 W8A8 against torch reference.""" + set_random_seed(seed) + + a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5) + w1_q, w2_q, w1_s, w2_s = _make_int8_moe_weights(E, N, K) + + score = torch.randn(M, E, dtype=torch.bfloat16) + score = torch.softmax(score, dim=-1, dtype=torch.float32) + topk_weight, topk_ids = torch.topk(score, topk) + topk_ids = topk_ids.to(torch.int32) + + ref_out = _ref_int8_moe(a, w1_q, w2_q, w1_s, w2_s, topk_weight, topk_ids) + + w1 = _prepack_experts(w1_q) if is_vnni else w1_q + w2 = _prepack_experts(w2_q) if is_vnni else w2_q + + out = ops.fused_experts_cpu( + a.clone(), + w1, + w2, + topk_weight, + topk_ids, + inplace, + ops.CPUQuantMethod.INT8_W8A8, + w1_s, + w2_s, + None, # w1_zero + None, # w2_zero + None, # block_size + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + is_vnni, + ) + torch.testing.assert_close( + ref_out.bfloat16(), + out, + atol=2e-1, + rtol=2e-1, + ) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/quantization/test_cpu_w8a8.py b/tests/quantization/test_cpu_w8a8.py new file mode 100644 index 00000000000..457aba2c6de --- /dev/null +++ b/tests/quantization/test_cpu_w8a8.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from vllm.platforms import current_platform + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +MODELS = [ + "RedHatAI/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8", # INT8 W8A8 MoE +] +DTYPE = ["bfloat16"] + + +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("dtype", DTYPE) +def test_cpu_w8a8(vllm_runner, model, dtype): + with vllm_runner(model, dtype=dtype) as llm: + output = llm.generate_greedy(["The capital of France is"], max_tokens=32) + assert output + print(output) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index cd67207b710..3ed1734cb91 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU FP8 W8A16 and MXFP4 W4A16 fused MoE experts.""" +"""CPU quantized fused MoE experts.""" import torch @@ -23,10 +23,16 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, kFp8Static128BlockSym, kInt4Static, + kInt8DynamicTokenSym, + kInt8StaticChannelSym, kMxfp4Static, ) from vllm.platforms import current_platform +# =========================================================================== +# FP8 W8A16 MoE +# =========================================================================== + def prepare_fp8_moe_layer_for_cpu( w13: torch.Tensor, @@ -177,6 +183,11 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): ) +# =========================================================================== +# MXFP4 W4A16 MoE +# =========================================================================== + + def prepare_mxfp4_moe_layer_for_cpu( w13: torch.Tensor, w2: torch.Tensor, @@ -326,6 +337,11 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): ) +# =========================================================================== +# INT4 W4A16 MoE +# =========================================================================== + + def prepare_int4_moe_layer_for_cpu( w13_packed: torch.Tensor, w2_packed: torch.Tensor, @@ -529,3 +545,159 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): None, # limit True, # is_vnni ) + + +# =========================================================================== +# INT8 W8A8 MoE +# =========================================================================== + + +def prepare_int8_moe_layer_for_cpu( + w13: torch.Tensor, + w2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """VNNI-prepack INT8 MoE weights for CPU kernel.""" + packed_w13 = torch.ops._C.convert_weight_packed(w13) + packed_w2 = torch.ops._C.convert_weight_packed(w2) + return packed_w13, packed_w2 + + +class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): + """CPU INT8 W8A8 per-channel weight / dynamic per-token activation + monolithic MoE experts.""" + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__( + moe_config, + quant_config, + ) + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cpu() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """VNNI-prepack INT8 MoE weights for CPU kernel.""" + from vllm.model_executor.utils import replace_parameter + + w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + select_experts, + ) + + topk_weights, topk_ids = select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + use_grouped_topk=num_expert_group is not None, + top_k=self.moe_config.experts_per_token, + renormalize=self.moe_config.routing_method + in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ), + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func="softmax", + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + e_score_correction_bias=e_score_correction_bias, + ) + + return fused_experts_cpu( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + False, # inplace + CPUQuantMethod.INT8_W8A8, + self.w1_scale, + self.w2_scale, + None, # w1_zero + None, # w2_zero + None, # block_size (per-channel, no block) + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + True, # is_vnni + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 6d50a3ba0ee..e31a3ca07ee 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -22,12 +22,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8DynamicTokenSym, kInt8StaticChannelSym, ) +from vllm.platforms import current_platform logger = init_logger(__name__) class Int8MoeBackend(Enum): TRITON = "TRITON" + CPU = "CPU" def _get_priority_backends( @@ -36,7 +38,18 @@ def _get_priority_backends( """ Get available backends in priority order based on platform and config. """ - return [Int8MoeBackend.TRITON] + _AVAILABLE_BACKENDS = [ + Int8MoeBackend.TRITON, + Int8MoeBackend.CPU, + ] + + def _move_to_front(backends: list[Int8MoeBackend], backend: Int8MoeBackend) -> None: + backends.insert(0, backends.pop(backends.index(backend))) + + if current_platform.is_cpu(): + _move_to_front(_AVAILABLE_BACKENDS, Int8MoeBackend.CPU) + + return _AVAILABLE_BACKENDS def backend_to_kernel_cls( @@ -49,6 +62,13 @@ def backend_to_kernel_cls( return [TritonExperts] + elif backend == Int8MoeBackend.CPU: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUExpertsInt8, + ) + + return [CPUExpertsInt8] + else: raise ValueError(f"Unknown Int8 MoE backend: {backend.value}") @@ -176,6 +196,24 @@ def make_int8_moe_quant_config( ) +def convert_to_int8_moe_kernel_format( + int8_backend: Int8MoeBackend, + w13: torch.Tensor, + w2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert INT8 MoE weights to backend-specific kernel format.""" + if int8_backend == Int8MoeBackend.CPU: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + prepare_int8_moe_layer_for_cpu, + ) + + w13, w2 = prepare_int8_moe_layer_for_cpu(w13, w2) + elif int8_backend != Int8MoeBackend.TRITON: + raise ValueError(f"Unsupported Int8 MoE backend: {int8_backend.value}") + + return w13, w2 + + def make_int8_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py index 74bf8a3546e..c29472cfc6b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - import torch from compressed_tensors.quantization import ( QuantizationArgs, @@ -20,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + convert_to_int8_moe_kernel_format, make_int8_moe_kernel, make_int8_moe_quant_config, select_int8_moe_backend, @@ -31,7 +31,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8DynamicTokenSym, kInt8StaticChannelSym, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.utils import replace_parameter, set_weight_attrs logger = init_logger(__name__) @@ -142,6 +142,14 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): layer.w2_input_scale = None def process_weights_after_loading(self, layer: RoutedExperts) -> None: + w13, w2 = convert_to_int8_moe_kernel_format( + int8_backend=self.int8_backend, + w13=layer.w13_weight, + w2=layer.w2_weight, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None self.moe_kernel = make_int8_moe_kernel( @@ -193,3 +201,27 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> 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, + ) From 5274c1181dc61bdf6e5eb610d37ebef694b1340d Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 29 Jun 2026 11:39:04 +0800 Subject: [PATCH 102/138] [Rust Frontend] Add Harmony Renderer for GPT-OSS (#46800) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 + rust/src/chat/Cargo.toml | 1 + rust/src/chat/src/backend/hf.rs | 56 +- rust/src/chat/src/lib.rs | 4 +- rust/src/chat/src/output/harmony/mod.rs | 21 +- .../src/chat/src/renderer/harmony/encoding.rs | 21 + .../harmony/fixtures/assistant_history.json | 14 + .../harmony/fixtures/assistant_history.txt | 7 + .../harmony/fixtures/developer_tools.json | 27 + .../harmony/fixtures/developer_tools.txt | 23 + .../harmony/fixtures/drop_analysis.json | 15 + .../harmony/fixtures/drop_analysis.txt | 7 + .../harmony/fixtures/leading_system.json | 13 + .../harmony/fixtures/leading_system.txt | 9 + .../harmony/fixtures/request_tools.json | 26 + .../harmony/fixtures/request_tools.txt | 19 + .../harmony/fixtures/simple_user.json | 6 + .../renderer/harmony/fixtures/simple_user.txt | 7 + .../fixtures/system_instructions_env.txt | 8 + .../harmony/fixtures/tool_roundtrip.json | 43 ++ .../harmony/fixtures/tool_roundtrip.txt | 19 + rust/src/chat/src/renderer/harmony/mod.rs | 487 ++++++++++++++++++ rust/src/chat/src/renderer/harmony/tests.rs | 212 ++++++++ rust/src/chat/src/renderer/mod.rs | 2 + rust/src/chat/src/renderer/selection.rs | 10 +- rust/src/chat/src/renderer/test_utils.rs | 5 +- rust/src/cmd/src/cli/tests.rs | 2 +- 27 files changed, 1033 insertions(+), 32 deletions(-) create mode 100644 rust/src/chat/src/renderer/harmony/encoding.rs create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/leading_system.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/request_tools.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/simple_user.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json create mode 100644 rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt create mode 100644 rust/src/chat/src/renderer/harmony/mod.rs create mode 100644 rust/src/chat/src/renderer/harmony/tests.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e1c051a6fe1..7820a7b6767 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5049,6 +5049,7 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "thiserror-ext", + "time", "tokio", "tracing", "tracing-subscriber", diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index cb28b1e9c14..40498bac3fe 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -25,6 +25,7 @@ strum.workspace = true subenum.workspace = true thiserror.workspace = true thiserror-ext.workspace = true +time.workspace = true tokio.workspace = true tracing.workspace = true trait-set.workspace = true diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 77ed24de854..9dff25ea49b 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -15,7 +15,9 @@ use crate::output::{ DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides, }; use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; -use crate::renderer::{DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer}; +use crate::renderer::{ + DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, +}; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -61,6 +63,7 @@ impl HfChatBackend { )?), RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()), RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), + RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), }; info!( @@ -148,13 +151,15 @@ mod tests { use std::sync::Arc; use tempfile::tempdir; + use thiserror_ext::AsReport as _; + use vllm_text::Prompt; use vllm_text::backend::hf::TokenizerSource; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; use super::HfChatBackend; - use crate::RendererSelection; - use crate::backend::{ChatBackend, LoadModelBackendsOptions}; + use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions}; use crate::request::{ChatContent, ChatMessage, ChatRequest}; + use crate::{ParserSelection, RendererSelection}; fn request_with_user_text(text: &str) -> ChatRequest { ChatRequest { @@ -219,12 +224,12 @@ mod tests { Arc::new(TestTokenizer) } - fn render_prompt( + fn backend_for_selection( renderer: RendererSelection, config_json: &str, tokenizer_config_json: &str, - ) -> String { - let backend = HfChatBackend::from_resolved_model_files( + ) -> HfChatBackend { + HfChatBackend::from_resolved_model_files( resolved_files(config_json, tokenizer_config_json), "test-model".to_string(), LoadModelBackendsOptions { @@ -236,9 +241,15 @@ mod tests { }, test_tokenizer(), ) - .unwrap(); + .unwrap() + } - backend + fn render_prompt( + renderer: RendererSelection, + config_json: &str, + tokenizer_config_json: &str, + ) -> String { + backend_for_selection(renderer, config_json, tokenizer_config_json) .chat_renderer() .render(&request_with_user_text("hello")) .unwrap() @@ -272,6 +283,35 @@ mod tests { assert_eq!(prompt, "hello"); } + #[test] + fn auto_uses_harmony_renderer_and_output_processor_for_gpt_oss_model_type() { + let backend = backend_for_selection( + RendererSelection::Auto, + r#"{"model_type":"gpt_oss"}"#, + r#"{"chat_template":"{{ messages[0].content }}"}"#, + ); + + let prompt = + backend.chat_renderer().render(&request_with_user_text("hello")).unwrap().prompt; + assert!(matches!(prompt, Prompt::TokenIds(_))); + + let mut request = request_with_user_text("hello"); + let error = match backend.new_chat_output_processor( + &mut request, + NewChatOutputProcessorOptions { + tool_call_parser: &ParserSelection::Explicit("json".to_string()), + reasoning_parser: &ParserSelection::Auto, + }, + ) { + Ok(_) => panic!("gpt_oss should reject generic parser overrides"), + Err(error) => error, + }; + assert_eq!( + error.to_report_string(), + "gpt_oss uses native Harmony output parsing; generic tool parser override `json` is not supported" + ); + } + #[test] fn language_model_only_skips_multimodal_preprocessor_config() { let mut files = resolved_files( diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 06e46f64b85..c16921ea758 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -29,8 +29,8 @@ pub use parser::reasoning::{ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ - ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, RenderedPrompt, - RendererSelection, + ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, + HarmonyChatRenderer, RenderedPrompt, RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index 4209dc0735c..597e3133795 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -4,16 +4,10 @@ //! `DecodedTextEvent` token IDs directly and lets the official `openai-harmony` //! parser recover the structured assistant message shape at token granularity. -use std::sync::LazyLock; - -use anyhow::Context; use asynk_strim_attr::{TryYielder, try_stream}; use futures::StreamExt as _; use openai_harmony::chat::{Content as HarmonyContent, Message as HarmonyMessage, Role}; -use openai_harmony::{ - HarmonyEncoding, HarmonyEncodingName, StreamableParser, load_harmony_encoding, -}; -use thiserror_ext::AsReport; +use openai_harmony::{HarmonyEncoding, StreamableParser}; use vllm_text::output::DecodedTextEvent; use crate::Result as ChatResult; @@ -24,6 +18,7 @@ use crate::output::{ generate_tool_call_id, }; use crate::parser::ParserSelection; +use crate::renderer::harmony::encoding::harmony_encoding; use crate::request::ChatRequest; /// Request-scoped Harmony output processor used for `model_type == "gpt_oss"`. @@ -384,18 +379,6 @@ async fn harmony_assistant_event_stream( Ok(()) } -/// Lazily load the shared GPT-OSS Harmony encoding once per process. -fn harmony_encoding() -> Result<&'static HarmonyEncoding> { - static ENCODING: LazyLock> = LazyLock::new(|| { - load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss) - .context("failed to load harmony encoding for gpt-oss") - }); - - ENCODING.as_ref().map_err(|error| Error::HarmonyOutputParsing { - error: error.to_report_string().into(), - }) -} - fn harmony_output_parsing_error( error: impl Into>, ) -> Error { diff --git a/rust/src/chat/src/renderer/harmony/encoding.rs b/rust/src/chat/src/renderer/harmony/encoding.rs new file mode 100644 index 00000000000..3b8030292d6 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/encoding.rs @@ -0,0 +1,21 @@ +//! Shared Harmony encoding helper for the GPT-OSS renderer and output parser. + +use std::sync::LazyLock; + +use anyhow::Context as _; +use openai_harmony::{HarmonyEncoding, HarmonyEncodingName, load_harmony_encoding}; +use thiserror_ext::AsReport as _; + +use crate::error::{Error, Result}; + +/// Lazily load the shared GPT-OSS Harmony encoding once per process. +pub(crate) fn harmony_encoding() -> Result<&'static HarmonyEncoding> { + static ENCODING: LazyLock> = LazyLock::new(|| { + load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss) + .context("failed to load harmony encoding for gpt-oss") + }); + + ENCODING.as_ref().map_err(|error| Error::HarmonyOutputParsing { + error: error.to_report_string().into(), + }) +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json new file mode 100644 index 00000000000..50edd03ee42 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json @@ -0,0 +1,14 @@ +{ + "add_generation_prompt": false, + "messages": [ + { + "role": "user", + "content": "What is 2 + 2?" + }, + { + "role": "assistant", + "reasoning_content": "Need simple arithmetic.", + "content": "4" + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt new file mode 100644 index 00000000000..dc08897e228 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt @@ -0,0 +1,7 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant<|channel|>final<|message|>4<|end|> diff --git a/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json new file mode 100644 index 00000000000..4516e3b32ba --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json @@ -0,0 +1,27 @@ +[ + { + "role": "developer", + "content": "Use tools when needed.", + "tools": [ + { + "function": { + "name": "lookup", + "description": "Lookup a record.", + "parameters": { + "type": "object", + "properties": { + "id": { + "type": "string" + } + }, + "required": ["id"] + } + } + } + ] + }, + { + "role": "user", + "content": "Find record abc." + } +] diff --git a/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt new file mode 100644 index 00000000000..a63b447f29c --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt @@ -0,0 +1,23 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message. +Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Instructions + +Use tools when needed. + +# Tools + +## functions + +namespace functions { + +// Lookup a record. +type lookup = (_: { +id: string, +}) => any; + +} // namespace functions<|end|><|start|>user<|message|>Find record abc.<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json new file mode 100644 index 00000000000..75fd5d8b29e --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json @@ -0,0 +1,15 @@ +[ + { + "role": "user", + "content": "What is 2 + 2?" + }, + { + "role": "assistant", + "reasoning_content": "This should be dropped.", + "content": "4" + }, + { + "role": "user", + "content": "What is 3 + 5?" + } +] diff --git a/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt new file mode 100644 index 00000000000..9e967b79564 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt @@ -0,0 +1,7 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant<|channel|>final<|message|>4<|end|><|start|>user<|message|>What is 3 + 5?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/leading_system.json b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.json new file mode 100644 index 00000000000..5ff190d0b85 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.json @@ -0,0 +1,13 @@ +{ + "reasoning_effort": "high", + "messages": [ + { + "role": "system", + "content": "Answer tersely." + }, + { + "role": "user", + "content": "What is 2 + 2?" + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt new file mode 100644 index 00000000000..e656a0a0a47 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt @@ -0,0 +1,9 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: high + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions + +Answer tersely.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/request_tools.json b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.json new file mode 100644 index 00000000000..db5988182fe --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.json @@ -0,0 +1,26 @@ +{ + "tools": [ + { + "function": { + "name": "get_weather", + "description": "Get weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + } + }, + "required": ["city"] + }, + "strict": true + } + } + ], + "messages": [ + { + "role": "user", + "content": "Check Hangzhou weather." + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt new file mode 100644 index 00000000000..f31bc449bfe --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt @@ -0,0 +1,19 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message. +Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Tools + +## functions + +namespace functions { + +// Get weather for a city. +type get_weather = (_: { +city: string, +}) => any; + +} // namespace functions<|end|><|start|>user<|message|>Check Hangzhou weather.<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/simple_user.json b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.json new file mode 100644 index 00000000000..b8b7f597d6f --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.json @@ -0,0 +1,6 @@ +[ + { + "role": "user", + "content": "Hello, who are you?" + } +] diff --git a/rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt new file mode 100644 index 00000000000..7e44ca314ce --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt @@ -0,0 +1,7 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>Hello, who are you?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt b/rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt new file mode 100644 index 00000000000..8ad0ac7d0ee --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt @@ -0,0 +1,8 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Answer tersely. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: high + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json new file mode 100644 index 00000000000..00ccae641e0 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json @@ -0,0 +1,43 @@ +{ + "tools": [ + { + "function": { + "name": "get_weather", + "description": "Get weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + } + }, + "required": ["city"] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "Check Hangzhou weather." + }, + { + "role": "assistant", + "reasoning_content": "Need current weather.", + "tool_calls": [ + { + "id": "call-weather", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Hangzhou\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call-weather", + "content": "{\"temperature\":20}" + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt new file mode 100644 index 00000000000..0e06a4d107e --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt @@ -0,0 +1,19 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message. +Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Tools + +## functions + +namespace functions { + +// Get weather for a city. +type get_weather = (_: { +city: string, +}) => any; + +} // namespace functions<|end|><|start|>user<|message|>Check Hangzhou weather.<|end|><|start|>assistant<|channel|>analysis<|message|>Need current weather.<|end|><|start|>assistant<|channel|>commentary to=functions.get_weather <|constrain|>json<|message|>{"city":"Hangzhou"}<|call|><|start|>functions.get_weather<|channel|>commentary to=assistant<|message|>{"temperature":20}<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs new file mode 100644 index 00000000000..70a1bb063e3 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -0,0 +1,487 @@ +//! Native Harmony chat renderer for `gpt_oss`. + +pub(crate) mod encoding; + +use openai_harmony::HarmonyEncoding; +use openai_harmony::chat::{ + Author, Conversation, DeveloperContent, Message, ReasoningEffort as HarmonyReasoningEffort, + Role, SystemContent, ToolDescription, +}; +use thiserror_ext::AsReport as _; +use time::macros::format_description; +use vllm_text::Prompt; + +use self::encoding::harmony_encoding; +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::error::{Error, Result}; +use crate::event::AssistantContentBlock; +use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatTool, GenerationPromptMode}; +use crate::{AssistantMessageExt as _, ReasoningEffort}; + +const SYSTEM_START_DATE_ENV: &str = "VLLM_SYSTEM_START_DATE"; +const HARMONY_SYSTEM_INSTRUCTIONS_ENV: &str = "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS"; + +/// GPT-OSS renderer backed by the official Harmony encoding. +pub struct HarmonyChatRenderer { + encoding: &'static HarmonyEncoding, + options: Options, +} + +struct Options { + system_start_date: String, + use_system_instructions: bool, +} + +impl HarmonyChatRenderer { + /// Create a Harmony renderer for production use. + /// + /// Environment-derived options are resolved once at construction time: + /// + /// - `VLLM_SYSTEM_START_DATE` pins the Harmony system start date. When it is + /// unset, the renderer uses the current local date with a UTC fallback. + /// - `VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS` moves leading instructions + /// into the system model identity when set to a non-zero integer. + pub fn new() -> Result { + Self::with_options( + env_system_start_date(), + env_use_harmony_system_instructions(), + ) + } + + /// Create a Harmony renderer with explicit preamble options. + /// + /// Tests use this constructor to avoid process-global environment mutation. + /// Production code should call [`Self::new`] so the renderer observes the + /// same environment contract as the Python Harmony path. + pub fn with_options( + system_start_date: impl Into, + use_system_instructions: bool, + ) -> Result { + Ok(Self { + encoding: harmony_encoding()?, + options: Options { + system_start_date: system_start_date.into(), + use_system_instructions, + }, + }) + } + + /// Render a chat request directly to Harmony token IDs. + /// + /// Harmony owns both prompt formatting and tokenization, so the Rust + /// frontend bypasses the generic HF tokenizer path for GPT-OSS input. + fn render_token_ids(&self, request: &ChatRequest) -> Result> { + if request.has_multimodal() { + return Err(Error::UnsupportedMultimodalContent("image_url")); + } + if matches!( + request.chat_options.generation_prompt_mode, + GenerationPromptMode::ContinueFinalAssistant + ) { + return Err(Error::ChatTemplate( + "Harmony renderer does not support continue_final_message".to_string(), + )); + } + + let messages = auto_drop_analysis_messages(to_harmony_messages(request, &self.options)?); + let conversation = Conversation::from_messages(messages); + // Pass `None` so oss-harmony does not apply its narrower built-in + // analysis-drop policy after the Rust-side Python-parity cleanup above. + let token_ids = match request.chat_options.generation_prompt_mode { + GenerationPromptMode::StartNewAssistant => self + .encoding + .render_conversation_for_completion(&conversation, Role::Assistant, None), + GenerationPromptMode::NoGenerationPrompt => { + self.encoding.render_conversation(&conversation, None) + } + GenerationPromptMode::ContinueFinalAssistant => unreachable!("checked above"), + } + .map_err(|error| { + Error::ChatTemplate(format!( + "failed to render Harmony prompt: {}", + error.as_report() + )) + })?; + + Ok(token_ids) + } +} + +impl ChatRenderer for HarmonyChatRenderer { + /// Render a chat request as [`Prompt::TokenIds`] with template kwargs echoed + /// for downstream accounting/debugging. + fn render(&self, request: &ChatRequest) -> Result { + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(self.render_token_ids(request)?), + effective_template_kwargs: request_template_kwargs(request), + }) + } +} + +/// Convert a vLLM chat request into a full Harmony conversation. +/// +/// This adds the Harmony system/developer preamble, peels at most one leading +/// system/developer instruction message, and then lowers the remaining chat +/// history message-by-message. +fn to_harmony_messages(request: &ChatRequest, options: &Options) -> Result> { + let (instructions, leading_developer_tools, remaining_messages) = + peel_leading_instructions(&request.messages)?; + let tool_call_names = tool_call_names(&request.messages); + let mut messages = + build_harmony_preamble(request, instructions, leading_developer_tools, options)?; + + for message in remaining_messages { + messages.extend(to_harmony_message(message, &tool_call_names, options)?); + } + + Ok(messages) +} + +/// Extract the optional leading instruction message used by the Harmony preamble. +/// +/// Python only peels the first leading `system` or `developer` message. Later +/// system/developer messages stay in the conversation and are lowered normally. +#[allow(clippy::type_complexity)] +fn peel_leading_instructions( + messages: &[ChatMessage], +) -> Result<(Option, Option<&[ChatTool]>, &[ChatMessage])> { + let Some(first) = messages.first() else { + return Ok((None, None, messages)); + }; + + match first { + ChatMessage::System { content } => Ok((Some(flatten_text(content)?), None, &messages[1..])), + ChatMessage::Developer { content, tools } => Ok(( + Some(flatten_text(content)?), + tools.as_deref(), + &messages[1..], + )), + ChatMessage::User { .. } + | ChatMessage::Assistant { .. } + | ChatMessage::ToolResponse { .. } => Ok((None, None, messages)), + } +} + +/// Build the Harmony preamble for one request. +/// +/// The preamble always contains a system message with date and reasoning-effort +/// metadata. Leading instructions live either in the system model identity or in +/// a developer message depending on `use_system_instructions`; request-level and +/// leading developer tools are attached to the developer message. +fn build_harmony_preamble( + request: &ChatRequest, + instructions: Option, + leading_developer_tools: Option<&[ChatTool]>, + options: &Options, +) -> Result> { + let mut messages = vec![Message::from_role_and_content( + Role::System, + system_content( + instructions.as_deref().filter(|_| options.use_system_instructions), + request.chat_options.reasoning_effort, + &options.system_start_date, + )?, + )]; + + let mut developer = DeveloperContent::new(); + let mut has_developer_content = false; + + if !options.use_system_instructions + && let Some(instructions) = instructions.as_deref().filter(|text| !text.is_empty()) + { + developer = developer.with_instructions(instructions); + has_developer_content = true; + } + + let tool_descriptions = preamble_tool_descriptions(request, leading_developer_tools); + if !tool_descriptions.is_empty() { + developer = developer.with_function_tools(tool_descriptions); + has_developer_content = true; + } + + if has_developer_content { + messages.push(Message::from_role_and_content(Role::Developer, developer)); + } + + Ok(messages) +} + +/// Collect request-level and leading developer function tools for the preamble. +fn preamble_tool_descriptions( + request: &ChatRequest, + leading_developer_tools: Option<&[ChatTool]>, +) -> Vec { + let mut tools = Vec::new(); + if request.tool_parsing_enabled() { + tools.extend(to_tool_descriptions(&request.tools)); + } + if let Some(leading_developer_tools) = leading_developer_tools { + tools.extend(to_tool_descriptions(leading_developer_tools)); + } + tools +} + +/// Construct the Harmony system content for the request preamble. +/// +/// Harmony defaults the reasoning effort to `medium` when none is provided, so +/// this only sets an explicit effort after validating vLLM's request value. +fn system_content( + instructions: Option<&str>, + reasoning_effort: Option, + system_start_date: &str, +) -> Result { + let mut content = + SystemContent::new().with_conversation_start_date(system_start_date.to_string()); + + if let Some(reasoning_effort) = reasoning_effort { + content = content.with_reasoning_effort(to_harmony_reasoning_effort(reasoning_effort)?); + } + + if let Some(instructions) = instructions.filter(|text| !text.is_empty()) { + let model_identity = match content.model_identity.as_deref() { + Some(identity) if !identity.is_empty() => format!("{identity}\n{instructions}"), + _ => instructions.to_string(), + }; + content = content.with_model_identity(model_identity); + } + + Ok(content) +} + +/// Lower a single vLLM chat message into one or more Harmony messages. +/// +/// Assistant messages can split into separate analysis, final, commentary, and +/// tool-call messages. Tool responses require the earlier assistant tool-call ID +/// map so the Harmony tool author can include `functions.{name}`. +fn to_harmony_message( + message: &ChatMessage, + tool_call_names: &std::collections::HashMap, + options: &Options, +) -> Result> { + Ok(match message { + ChatMessage::System { content } => { + let instructions = flatten_text(content)?; + vec![system_or_developer_message( + "system", + instructions, + None, + options, + )?] + } + ChatMessage::Developer { content, tools } => { + let instructions = flatten_text(content)?; + vec![developer_message(Some(instructions), tools.as_deref())] + } + ChatMessage::User { content } => { + vec![Message::from_role_and_content( + Role::User, + flatten_text(content)?, + )] + } + ChatMessage::Assistant { content } => assistant_messages(content), + ChatMessage::ToolResponse { + content, + tool_call_id, + } => { + let name = tool_call_names.get(tool_call_id).ok_or_else(|| { + Error::ChatTemplate(format!( + "invalid Harmony tool message: unknown tool_call_id `{tool_call_id}`" + )) + })?; + vec![ + Message::from_author_and_content( + Author::new(Role::Tool, format!("functions.{name}")), + flatten_text(content)?, + ) + .with_channel("commentary") + .with_recipient("assistant"), + ] + } + }) +} + +/// Lower a non-leading system/developer message. +/// +/// Harmony treats most extra system/developer messages as developer +/// instructions. When system-instructions mode is enabled, system messages are +/// rendered as system model-identity additions to match Python. +fn system_or_developer_message( + role: &str, + instructions: String, + tools: Option<&[ChatTool]>, + options: &Options, +) -> Result { + if role == "system" && options.use_system_instructions { + return Ok(Message::from_role_and_content( + Role::System, + system_content(Some(&instructions), None, &options.system_start_date)?, + )); + } + + Ok(developer_message(Some(instructions), tools)) +} + +/// Build a Harmony developer message with optional instructions and function tools. +fn developer_message(instructions: Option, tools: Option<&[ChatTool]>) -> Message { + let mut content = DeveloperContent::new(); + if let Some(instructions) = instructions.filter(|text| !text.is_empty()) { + content = content.with_instructions(instructions); + } + if let Some(tools) = tools { + let tools = to_tool_descriptions(tools); + if !tools.is_empty() { + content = content.with_function_tools(tools); + } + } + Message::from_role_and_content(Role::Developer, content) +} + +/// Lower assistant history into Harmony channels. +/// +/// Plain assistant text goes to `final`. When the assistant has tool calls, +/// visible text goes to `commentary`, reasoning goes to `analysis`, and each +/// function call becomes a `commentary` message to `functions.{name}` with JSON +/// constrained content. +fn assistant_messages(content: &[AssistantContentBlock]) -> Vec { + let mut messages = Vec::new(); + let has_tool_calls = content.has_tool_calls(); + + if has_tool_calls { + let text = content.text(); + if !text.is_empty() { + messages.push( + Message::from_role_and_content(Role::Assistant, text).with_channel("commentary"), + ); + } + } + + if let Some(reasoning) = content.reasoning() { + messages.push( + Message::from_role_and_content(Role::Assistant, reasoning).with_channel("analysis"), + ); + } + + if has_tool_calls { + for tool_call in content.tool_calls() { + messages.push( + Message::from_role_and_content(Role::Assistant, tool_call.arguments.clone()) + .with_channel("commentary") + .with_recipient(format!("functions.{}", tool_call.name)) + .with_content_type("<|constrain|>json"), + ); + } + } else { + let text = content.text(); + if !text.is_empty() { + messages + .push(Message::from_role_and_content(Role::Assistant, text).with_channel("final")); + } + } + + messages +} + +/// Build the tool-call ID to function-name map used by later tool responses. +fn tool_call_names(messages: &[ChatMessage]) -> std::collections::HashMap { + let mut names = std::collections::HashMap::new(); + for message in messages { + let ChatMessage::Assistant { content } = message else { + continue; + }; + for tool_call in content.tool_calls() { + names.insert(tool_call.id.clone(), tool_call.name.clone()); + } + } + names +} + +/// Drop stale assistant analysis messages using vLLM Python's policy. +/// +/// Once an assistant final message exists, earlier analysis messages represent +/// chain-of-thought for completed turns and should not be replayed to the model. +fn auto_drop_analysis_messages(messages: Vec) -> Vec { + // Match vLLM Python's Harmony cleanup: once an assistant final message exists, + // previous assistant analysis messages are stale chain-of-thought and should + // be removed. oss-harmony can also drop analysis with `Some(Default::default())`, + // but that built-in path only triggers when the last assistant message is final + // and drops relative to the first final message, which misses longer multi-turn + // histories with later user/tool turns. + let Some(last_assistant_final_index) = messages.iter().rposition(|message| { + message.author.role == Role::Assistant && message.channel.as_deref() == Some("final") + }) else { + return messages; + }; + + messages + .into_iter() + .enumerate() + .filter_map(|(index, message)| { + (index >= last_assistant_final_index || message.channel.as_deref() != Some("analysis")) + .then_some(message) + }) + .collect() +} + +/// Flatten vLLM text content and reject unsupported multimodal parts. +fn flatten_text(content: &ChatContent) -> Result { + content.try_flatten_to_text() +} + +/// Convert vLLM function tool definitions to Harmony tool descriptions. +fn to_tool_descriptions(tools: &[ChatTool]) -> Vec { + tools + .iter() + .map(|tool| { + ToolDescription::new( + tool.name.clone(), + tool.description.clone().unwrap_or_default(), + Some(tool.parameters.clone()), + ) + }) + .collect() +} + +/// Map supported OpenAI reasoning-effort values onto Harmony's enum. +fn to_harmony_reasoning_effort( + reasoning_effort: ReasoningEffort, +) -> Result { + match reasoning_effort { + ReasoningEffort::Low => Ok(HarmonyReasoningEffort::Low), + ReasoningEffort::Medium => Ok(HarmonyReasoningEffort::Medium), + ReasoningEffort::High => Ok(HarmonyReasoningEffort::High), + ReasoningEffort::None + | ReasoningEffort::Minimal + | ReasoningEffort::XHigh + | ReasoningEffort::Max => Err(Error::ChatTemplate(format!( + "reasoning_effort={:?} is not supported by Harmony. Supported values are: low, medium, high.", + reasoning_effort.as_str() + ))), + } +} + +/// Resolve the system start date from the environment or the current date. +fn env_system_start_date() -> String { + std::env::var(SYSTEM_START_DATE_ENV) + .ok() + .filter(|date| !date.is_empty()) + .unwrap_or_else(current_date) +} + +/// Format today's date as `YYYY-MM-DD`, preferring local time. +fn current_date() -> String { + const DATE_FORMAT: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day]"); + let now = time::OffsetDateTime::now_local().unwrap_or_else(|_| time::OffsetDateTime::now_utc()); + now.format(DATE_FORMAT).expect("static date format should be valid") +} + +/// Resolve the env flag that places leading instructions in system identity. +fn env_use_harmony_system_instructions() -> bool { + std::env::var(HARMONY_SYSTEM_INSTRUCTIONS_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .is_some_and(|value| value != 0) +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs new file mode 100644 index 00000000000..bcbf97c8664 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -0,0 +1,212 @@ +use std::path::PathBuf; + +use expect_test::{ExpectFile, expect, expect_file}; +use thiserror_ext::AsReport as _; + +use super::HarmonyChatRenderer; +use super::encoding::harmony_encoding; +use crate::ChatRenderer; +use crate::error::Error; +use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ + ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, +}; + +const PINNED_DATE: &str = "2025-06-28"; + +fn fixture_request(input_name: &str) -> ChatRequest { + fixture_chat_request( + &fixture_path(input_name), + FixtureRequestOptions { + enable_thinking: false, + no_generation_prompt_when_last_assistant: false, + }, + ) +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/harmony") + .join("fixtures") + .join(name) +} + +fn test_renderer(use_system_instructions: bool) -> HarmonyChatRenderer { + HarmonyChatRenderer::with_options(PINNED_DATE, use_system_instructions).unwrap() +} + +fn render_token_ids(request: &ChatRequest) -> Vec { + render_token_ids_with(&test_renderer(false), request) +} + +fn render_token_ids_with(renderer: &HarmonyChatRenderer, request: &ChatRequest) -> Vec { + renderer + .render(request) + .unwrap() + .prompt + .into_token_ids() + .expect("Harmony renderer returns token IDs") +} + +fn render_prompt_text(request: &ChatRequest) -> String { + render_prompt_text_with(&test_renderer(false), request) +} + +fn render_prompt_text_with(renderer: &HarmonyChatRenderer, request: &ChatRequest) -> String { + let token_ids = render_token_ids_with(renderer, request); + harmony_encoding().unwrap().tokenizer().decode_utf8(&token_ids).unwrap() +} + +fn assert_fixture(input_name: &str, expected: ExpectFile) { + let request = fixture_request(input_name); + let rendered = format!("{}\n", render_prompt_text(&request)); + expected.assert_eq(&rendered); +} + +#[test] +fn renders_token_ids() { + let request = fixture_request("simple_user.json"); + + assert!(!render_token_ids(&request).is_empty()); +} + +#[test] +fn renders_simple_user_fixture() { + assert_fixture("simple_user.json", expect_file!["fixtures/simple_user.txt"]); +} + +#[test] +fn renders_leading_system_fixture() { + assert_fixture( + "leading_system.json", + expect_file!["fixtures/leading_system.txt"], + ); +} + +#[test] +fn renders_system_instructions_env_fixture() { + let renderer = test_renderer(true); + let request = fixture_request("leading_system.json"); + let rendered = format!("{}\n", render_prompt_text_with(&renderer, &request)); + expect_file!["fixtures/system_instructions_env.txt"].assert_eq(&rendered); +} + +#[test] +fn renders_request_tools_fixture() { + assert_fixture( + "request_tools.json", + expect_file!["fixtures/request_tools.txt"], + ); +} + +#[test] +fn renders_developer_tools_fixture() { + assert_fixture( + "developer_tools.json", + expect_file!["fixtures/developer_tools.txt"], + ); +} + +#[test] +fn renders_assistant_history_fixture() { + assert_fixture( + "assistant_history.json", + expect_file!["fixtures/assistant_history.txt"], + ); +} + +#[test] +fn renders_tool_roundtrip_fixture() { + assert_fixture( + "tool_roundtrip.json", + expect_file!["fixtures/tool_roundtrip.txt"], + ); +} + +#[test] +fn drops_stale_analysis_fixture() { + assert_fixture( + "drop_analysis.json", + expect_file!["fixtures/drop_analysis.txt"], + ); +} + +#[test] +fn rejects_invalid_reasoning_effort() { + let mut request = ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + + let error = test_renderer(false).render(&request).unwrap_err(); + + expect![[r#"chat template error: reasoning_effort="none" is not supported by Harmony. Supported values are: low, medium, high."#]] + .assert_eq(&error.to_report_string()); +} + +#[test] +fn rejects_unknown_tool_response_id() { + let request = ChatRequest { + messages: vec![ + ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall( + AssistantToolCall { + id: "call-known".to_string(), + name: "lookup".to_string(), + arguments: "{}".to_string(), + }, + )]), + ChatMessage::tool_response("{}", "call-unknown"), + ], + ..ChatRequest::for_test() + }; + + let error = test_renderer(false).render(&request).unwrap_err(); + + expect![ + "chat template error: invalid Harmony tool message: unknown tool_call_id `call-unknown`" + ] + .assert_eq(&error.to_report_string()); +} + +#[test] +fn rejects_multimodal_input() { + let request = ChatRequest { + messages: vec![ChatMessage::user(vec![ChatContentPart::image_url( + "data:image/png;base64,test", + )])], + ..ChatRequest::for_test() + }; + + let error = test_renderer(false).render(&request).unwrap_err(); + + assert!(matches!( + error, + Error::UnsupportedMultimodalContent("image_url") + )); +} + +#[test] +fn rejects_continue_final_assistant() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("write"), + ChatMessage::assistant_text("partial"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let error = test_renderer(false).render(&request).unwrap_err(); + + expect!["chat template error: Harmony renderer does not support continue_final_message"] + .assert_eq(&error.to_report_string()); +} + +#[test] +fn no_generation_prompt_omits_trailing_assistant_start() { + let mut request = fixture_request("simple_user.json"); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_prompt_text(&request); + + assert!(!rendered.ends_with("<|start|>assistant")); +} diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index c4ee787c868..f1c510a1b3d 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -9,6 +9,7 @@ use crate::request::{ChatRequest, ReasoningEffort}; pub mod deepseek_v32; pub mod deepseek_v4; +pub mod harmony; pub mod hf; mod selection; #[cfg(test)] @@ -16,6 +17,7 @@ mod test_utils; pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; +pub use harmony::HarmonyChatRenderer; pub use selection::RendererSelection; /// Rendered chat prompt submitted to the text backend. diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index 09bdd6b9721..837ec7d69c6 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -19,12 +19,16 @@ pub enum RendererSelection { DeepSeekV32, /// Force the DeepSeek V4 renderer. DeepSeekV4, + /// Force the GPT-OSS Harmony renderer. + Harmony, } impl RendererSelection { pub const AUTO_LITERAL: &str = "auto"; pub const DEEPSEEK_V32_LITERAL: &str = "deepseek_v32"; pub const DEEPSEEK_V4_LITERAL: &str = "deepseek_v4"; + pub const GPT_OSS_MODEL_TYPE: &str = "gpt_oss"; + pub const HARMONY_LITERAL: &str = "harmony"; pub const HF_LITERAL: &str = "hf"; /// Resolve the renderer selection using the given model type string, if @@ -34,6 +38,7 @@ impl RendererSelection { Self::Auto => match model_type { Self::DEEPSEEK_V32_LITERAL => Self::DeepSeekV32, Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4, + Self::GPT_OSS_MODEL_TYPE => Self::Harmony, _ => Self::Hf, }, selection => selection, @@ -53,6 +58,8 @@ impl FromStr for RendererSelection { Ok(Self::DeepSeekV32) } else if value.eq_ignore_ascii_case(Self::DEEPSEEK_V4_LITERAL) { Ok(Self::DeepSeekV4) + } else if value.eq_ignore_ascii_case(Self::HARMONY_LITERAL) { + Ok(Self::Harmony) } else { Err(format!( "unknown renderer `{value}` (expected one of: {})", @@ -69,6 +76,7 @@ impl fmt::Display for RendererSelection { Self::Hf => f.write_str(Self::HF_LITERAL), Self::DeepSeekV32 => f.write_str(Self::DEEPSEEK_V32_LITERAL), Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL), + Self::Harmony => f.write_str(Self::HARMONY_LITERAL), } } } @@ -95,7 +103,7 @@ mod tests { fn renderer_selection_expected_error_message() { let err = RendererSelection::from_str("unknown").unwrap_err(); expect_test::expect![ - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4)" + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony)" ] .assert_eq(&err); } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 0aab3769db4..bf560de8427 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -7,7 +7,7 @@ use serde_json::Value; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, - GenerationPromptMode, + GenerationPromptMode, ReasoningEffort, }; /// Options for constructing a [`ChatRequest`] from a fixture file. @@ -42,6 +42,7 @@ pub(crate) struct FixtureRequest { tools: Vec, messages: Vec, add_generation_prompt: Option, + reasoning_effort: Option, } impl FixtureFile { @@ -52,6 +53,7 @@ impl FixtureFile { tools: Vec::new(), messages, add_generation_prompt: None, + reasoning_effort: None, }, } } @@ -154,6 +156,7 @@ impl FixtureRequest { if self.add_generation_prompt == Some(false) { request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; } + request.chat_options.reasoning_effort = self.reasoning_effort; if options.enable_thinking { for key in ["thinking", "enable_thinking"] { request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true)); diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 345cc9f60d6..a9b11ca18f7 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -468,7 +468,7 @@ fn serve_args_reject_unknown_renderer_value() { .unwrap_err(); expect![[r#" - error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4) + error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony) For more information, try '--help'. "#]] From 4559c43a9526597c00cbcc4f59979496500268d1 Mon Sep 17 00:00:00 2001 From: Soyaazz <523420504@qq.com> Date: Mon, 29 Jun 2026 12:52:00 +0800 Subject: [PATCH 103/138] [MM][CG] Gemma3 Encoder CUDA Graph (#43591) Signed-off-by: JisoLya <523420504@qq.com> Signed-off-by: Soyaazz <523420504@qq.com> Co-authored-by: Isotr0py --- docs/design/cuda_graphs_multimodal.md | 1 + .../generation/test_vit_cudagraph.py | 14 ++ vllm/model_executor/models/gemma3_mm.py | 137 +++++++++++++++++- 3 files changed, 151 insertions(+), 1 deletion(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index ceefc195021..7eab425d6e2 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -127,6 +127,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | Architecture | Models | CG for Image | CG for Video | Dual-Path Graph | | ------------ | ------ | ------------ | ------------ | --------------- | | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | โœ…๏ธŽ | โŒ๏ธŽ | โœ…๏ธŽ | +| `Gemma3ForConditionalGeneration` | `Gemma3` | โœ…๏ธŽ | โŒ๏ธŽ | โŒ๏ธŽ | | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | โœ…๏ธŽ | โœ…๏ธŽ | โŒ๏ธŽ | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | โœ…๏ธŽ | โœ…๏ธŽ | โŒ๏ธŽ | | `KimiVLForConditionalGeneration` | `Kimi-VL` | โœ…๏ธŽ | โŒ๏ธŽ | โŒ๏ธŽ | diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 52b28ca8600..954bbdbb9b8 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -62,7 +62,21 @@ def step3_vl_chat_template(content: str) -> str: ) +def gemma3_chat_template(content: str) -> str: + return f"user\n{content}\nmodel\n" + + MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { + "gemma3": VitCudagraphTestConfig( + model="google/gemma-3-4b-it", + modalities=["image"], + image_prompt=gemma3_chat_template("What is in this image?"), + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [512], + }, + dtype="bfloat16", + max_model_len=4096, + ), "llama4": VitCudagraphTestConfig( model="meta-llama/Llama-4-Scout-17B-16E-Instruct", modalities=["image"], diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 6ecadbcd670..9e58438f4cf 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -39,6 +39,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( MultiModalEmbeddings, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -467,7 +468,7 @@ class Gemma3MultiModalProjector(nn.Module): dummy_inputs=Gemma3DummyInputsBuilder, ) class Gemma3ForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA + nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsEncoderCudaGraph ): packed_modules_mapping = { "qkv_proj": [ @@ -504,8 +505,12 @@ class Gemma3ForConditionalGeneration( quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config self.config = config + self.model_config = vllm_config.model_config self.quant_config = quant_config self.multimodal_config = multimodal_config + self.vit_positions_per_patch = ( + self.config.vision_config.image_size // self.config.vision_config.patch_size + ) ** 2 self.configure_mm_token_handling( vocab_size=config.text_config.vocab_size, @@ -682,3 +687,133 @@ class Gemma3ForConditionalGeneration( """ # The Gemma3 connector maintains a 1:1 token mapping return num_vision_tokens + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphConfig, + ) + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=["pixel_values"], + out_hidden_size=self.config.text_config.hidden_size, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = self.config.mm_tokens_per_image + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + num_patches = mm_kwargs["num_patches"] + mm_tokens_per_image = self.config.mm_tokens_per_image + + return [ + EncoderItemSpec( + input_size=int(np) * self.vit_positions_per_patch, + output_tokens=int(np) * mm_tokens_per_image, + ) + for np in num_patches + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + pixel_values = mm_kwargs["pixel_values"] + num_patches = mm_kwargs["num_patches"] + + if len(indices) == 0: + return { + "pixel_values": pixel_values[:0], + "num_patches": num_patches[:0], + } + cum_patches = [0] + for p in num_patches: + cum_patches.append(cum_patches[-1] + int(p)) + + selected_pv = torch.cat( + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices] + ) + selected_np = num_patches[indices] + + return { + "pixel_values": selected_pv, + "num_patches": selected_np, + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + mm_tokens_per_image = self.config.mm_tokens_per_image + num_images = min( + token_budget // mm_tokens_per_image, + max_batch_size, + ) + + image_size = self.config.vision_config.image_size + dummy_pixel_values = torch.randn( + num_images, + 3, + image_size, + image_size, + device=device, + dtype=dtype, + ) + values = {"pixel_values": dummy_pixel_values} + + return EncoderCudaGraphCaptureInputs( + values, + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + return EncoderCudaGraphReplayBuffers( + values={"pixel_values": mm_kwargs["pixel_values"]}, + ) + + def encoder_cudagraph_forward( + self, + values: dict[str, torch.Tensor], + ) -> torch.Tensor: + pixel_values = values["pixel_values"] + image_features = self.vision_tower(pixel_values) + image_features = self.multi_modal_projector(image_features) + return image_features.flatten(end_dim=1) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + ) -> torch.Tensor: + image_input = self._parse_and_validate_image_input(**mm_kwargs) + results = self._process_image_input(image_input) + return torch.cat(results, dim=0) From f6bb8682ee5b6a35cb0c74a4c1f01165ee6ca24d Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 29 Jun 2026 08:50:57 +0100 Subject: [PATCH 104/138] Fix docs on main (#47009) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/design/moe_kernel_features.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index d49790e833a..07d2a539801 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -89,7 +89,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],
[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] | | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | -| hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.experts.hpc.HPCExperts] | +| hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] | | rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | From db28ae2d078da82d01f8e7fad05fb27a52ce37ef Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 29 Jun 2026 02:59:24 -0500 Subject: [PATCH 105/138] [ROCm][CI] Explicitly tear down multimodal offline LLMs (#46999) Signed-off-by: Andreas Karatzas --- tests/conftest.py | 8 ++- tests/entrypoints/multimodal/conftest.py | 72 +++++++++++++++++++ tests/entrypoints/multimodal/llm/test_chat.py | 16 +---- .../llm/test_mm_cache_external_injection.py | 6 +- .../multimodal/llm/test_mm_cache_stats.py | 3 +- .../multimodal/llm/test_mm_embeds_only.py | 15 ++-- 6 files changed, 91 insertions(+), 29 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 4b92f285fac..6f9c8fa120f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1587,7 +1587,13 @@ class AssetHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", ctype) self.send_header("Content-Length", str(len(data))) self.end_headers() - self.wfile.write(data) + try: + self.wfile.write(data) + except (BrokenPipeError, ConnectionResetError) as e: + logger.debug( + "Client disconnected while serving test asset %s: %r", filename, e + ) + self.close_connection = True def _find_free_port() -> int: diff --git a/tests/entrypoints/multimodal/conftest.py b/tests/entrypoints/multimodal/conftest.py index 9c260bc2225..8003f1bf7dc 100644 --- a/tests/entrypoints/multimodal/conftest.py +++ b/tests/entrypoints/multimodal/conftest.py @@ -1,5 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Any + +import pytest # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) TEST_IMAGE_ASSETS = [ @@ -8,3 +13,70 @@ TEST_IMAGE_ASSETS = [ "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png", "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png", ] + + +def _shutdown_llm(llm: Any, gpu_memory_utilization: float) -> None: + from vllm.distributed import cleanup_dist_env_and_memory + from vllm.platforms import current_platform + + try: + shutdown_timeout = 60.0 if current_platform.is_rocm() else None + llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout) + except Exception: + pass + + del llm + + try: + import torch + + torch._dynamo.reset() + except Exception: + pass + + cleanup_dist_env_and_memory() + + if current_platform.is_rocm(): + from tests.utils import wait_for_rocm_memory_to_settle + + wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + + +@contextmanager +def managed_llm(*args: Any, **kwargs: Any) -> Iterator[Any]: + from vllm import LLM + + llm = LLM(*args, **kwargs) + gpu_memory_utilization = ( + llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization + ) + try: + yield llm + finally: + _shutdown_llm(llm, gpu_memory_utilization) + + +def _make_managed_llm_factory() -> Iterator[Callable[..., Any]]: + from vllm import LLM + + llms: list[tuple[Any, float]] = [] + + def make_llm(*args: Any, **kwargs: Any) -> Any: + llm = LLM(*args, **kwargs) + gpu_memory_utilization = ( + llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization + ) + llms.append((llm, gpu_memory_utilization)) + return llm + + try: + yield make_llm + finally: + while llms: + llm, gpu_memory_utilization = llms.pop() + _shutdown_llm(llm, gpu_memory_utilization) + + +@pytest.fixture +def multimodal_llm_factory() -> Iterator[Callable[..., Any]]: + yield from _make_managed_llm_factory() diff --git a/tests/entrypoints/multimodal/llm/test_chat.py b/tests/entrypoints/multimodal/llm/test_chat.py index b670c4c3c4e..4de1f5cb80a 100644 --- a/tests/entrypoints/multimodal/llm/test_chat.py +++ b/tests/entrypoints/multimodal/llm/test_chat.py @@ -1,19 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import weakref - import pytest from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory @pytest.fixture(scope="function") -def vision_llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( +def vision_llm(multimodal_llm_factory): + return multimodal_llm_factory( model="microsoft/Phi-3.5-vision-instruct", max_model_len=4096, max_num_seqs=5, @@ -23,12 +17,6 @@ def vision_llm(): seed=0, ) - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() - @pytest.mark.parametrize( "image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py index f3ae499d635..076a381f6cd 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py @@ -69,6 +69,7 @@ def test_inject_into_mm_cache( image_urls, mm_processor_cache_type, caplog_vllm, + multimodal_llm_factory, ): """Test that inject_into_mm_cache() injects pre-processed mm_kwargs into the processor cache and MM cache hit metrics are updated correctly. @@ -78,7 +79,7 @@ def test_inject_into_mm_cache( 2. Extract cached kwargs, call inject_into_mm_cache with a new hash, then generate with a pre-rendered input -> verifies injection works """ - llm = LLM( + llm = multimodal_llm_factory( model="llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, @@ -145,11 +146,12 @@ def test_inject_into_mm_cache( def test_inject_into_mm_cache_without_cache( num_gpus_available, image_urls, + multimodal_llm_factory, ): """Test that inject_into_mm_cache works gracefully when processor cache is disabled (mm_processor_cache_gb=0). Should not crash. """ - llm = LLM( + llm = multimodal_llm_factory( model="llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py index 496e98d5ca1..dbea37f64ee 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py @@ -61,8 +61,9 @@ def test_mm_cache_stats( image_urls, mm_processor_cache_type, caplog_vllm, + multimodal_llm_factory, ): - llm = LLM( + llm = multimodal_llm_factory( model="llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, diff --git a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py index 13d0fd58b13..57bec9c1188 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py +++ b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py @@ -1,13 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import weakref - import pytest +from tests.entrypoints.multimodal.conftest import managed_llm from vllm import LLM, SamplingParams from vllm.assets.image import ImageAsset -from vllm.distributed import cleanup_dist_env_and_memory MODEL = "llava-hf/llava-1.5-7b-hf" PROMPT = "USER: \nDescribe this image briefly.\nASSISTANT:" @@ -17,20 +15,15 @@ TEXT_ONLY_PROMPT = "USER: What is 2 + 2?\nASSISTANT:" @pytest.fixture(scope="module") def llm(): """LLM with enable_mm_embeds=True and all modality limits zeroed out.""" - llm = LLM( + with managed_llm( model=MODEL, max_model_len=2048, enforce_eager=True, gpu_memory_utilization=0.8, enable_mm_embeds=True, limit_mm_per_prompt={"image": 0}, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as llm: + yield llm @pytest.mark.skip_global_cleanup From 5051698e41b7dc3da421f1c50bfe178a92dc7881 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 29 Jun 2026 09:52:23 +0100 Subject: [PATCH 106/138] Remove unnecessary `load_weights` methods (#44589) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/model_executor/test_weight_utils.py | 121 +++++++++++++++++ vllm/lora/worker_manager.py | 6 +- vllm/model_executor/layers/linear.py | 122 ++++++++++++------ .../layers/quantization/base_config.py | 51 +++++++- .../model_executor/layers/quantization/fp8.py | 19 +-- .../layers/quantization/quark/quark.py | 19 +-- .../model_loader/bitsandbytes_loader.py | 3 +- .../model_loader/reload/layerwise.py | 11 +- vllm/model_executor/model_loader/utils.py | 2 +- vllm/model_executor/models/arcee.py | 77 ++--------- vllm/model_executor/models/chatglm.py | 53 +------- vllm/model_executor/models/cohere_eagle.py | 41 +----- vllm/model_executor/models/commandr.py | 66 ++-------- vllm/model_executor/models/exaone.py | 81 ++---------- vllm/model_executor/models/exaone4.py | 81 ++---------- vllm/model_executor/models/fairseq2_llama.py | 6 +- vllm/model_executor/models/gemma.py | 68 +++------- vllm/model_executor/models/gemma2.py | 71 ++-------- vllm/model_executor/models/gemma3.py | 75 ++--------- vllm/model_executor/models/glm4.py | 84 ++---------- vllm/model_executor/models/glm4v.py | 13 ++ vllm/model_executor/models/gpt_j.py | 62 ++------- vllm/model_executor/models/granite.py | 84 +++--------- vllm/model_executor/models/hyperclovax.py | 82 ++---------- vllm/model_executor/models/interfaces.py | 3 +- vllm/model_executor/models/internlm2.py | 52 ++------ vllm/model_executor/models/jais2.py | 69 ++-------- vllm/model_executor/models/jina.py | 5 +- vllm/model_executor/models/llama.py | 82 ++++-------- vllm/model_executor/models/mamba.py | 25 +--- vllm/model_executor/models/mamba2.py | 26 +--- vllm/model_executor/models/mimo.py | 57 ++------ vllm/model_executor/models/mistral_eagle.py | 10 +- vllm/model_executor/models/mpt.py | 17 --- vllm/model_executor/models/nemotron.py | 66 ++-------- vllm/model_executor/models/nemotron_nas.py | 70 ++-------- vllm/model_executor/models/olmo.py | 63 ++------- vllm/model_executor/models/olmo2.py | 62 ++------- vllm/model_executor/models/opt.py | 54 ++------ vllm/model_executor/models/orion.py | 53 ++------ vllm/model_executor/models/ouro.py | 75 ++--------- vllm/model_executor/models/phi.py | 60 ++------- vllm/model_executor/models/qwen2.py | 85 +++--------- vllm/model_executor/models/qwen2_rm.py | 12 +- vllm/model_executor/models/qwen3.py | 13 +- vllm/model_executor/models/rnj1.py | 86 ++---------- vllm/model_executor/models/seed_oss.py | 72 ++--------- vllm/model_executor/models/solar.py | 76 ++--------- vllm/model_executor/models/stablelm.py | 53 ++------ vllm/model_executor/models/starcoder2.py | 50 ++----- vllm/model_executor/models/step1.py | 61 +++------ .../models/transformers/base.py | 3 - vllm/model_executor/models/utils.py | 91 +++++++++---- vllm/model_executor/models/whisper.py | 47 ++----- 54 files changed, 821 insertions(+), 1975 deletions(-) diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 260ebdcefb3..9e67609b78e 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -160,5 +160,126 @@ class TestMaybeRemapKvScaleName: assert result is None +class TestKvCacheScaleMapper: + """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the + per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of + checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent + so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`.""" + + def _mapper(self): + # `get_cache_scale_mapper` does not use `self`; call it on the base + # class to get the default (non-config-specific) mapper. + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + + return QuantizationConfig.get_cache_scale_mapper() + + def _map(self, name: str) -> str | None: + return self._mapper()._map_name(name) + + @pytest.mark.parametrize( + "name,expected", + [ + # Qwen3-MoE / llm-compressor fused qkv_proj + ( + "model.layers.0.self_attn.qkv_proj.k_scale", + "model.layers.0.self_attn.attn.k_scale", + ), + ( + "model.layers.0.self_attn.qkv_proj.v_scale", + "model.layers.0.self_attn.attn.v_scale", + ), + # ModelOpt / NVFP4 k_proj/v_proj + ( + "model.layers.0.self_attn.k_proj.k_scale", + "model.layers.0.self_attn.attn.k_scale", + ), + ( + "model.layers.0.self_attn.v_proj.v_scale", + "model.layers.0.self_attn.attn.v_scale", + ), + # deprecated fused kv_scale and bare scales + ( + "model.layers.0.self_attn.kv_scale", + "model.layers.0.self_attn.attn.k_scale", + ), + ( + "model.layers.0.self_attn.k_scale", + "model.layers.0.self_attn.attn.k_scale", + ), + # NemotronH mixer + ( + "model.layers.0.mixer.k_proj.k_scale", + "model.layers.0.mixer.attn.k_scale", + ), + # already in vLLM form -> unchanged (idempotent) + ( + "model.layers.0.self_attn.attn.k_scale", + "model.layers.0.self_attn.attn.k_scale", + ), + # non-kv scales must not be touched + ( + "model.layers.0.self_attn.k_proj.weight_scale", + "model.layers.0.self_attn.k_proj.weight_scale", + ), + ( + "model.layers.0.self_attn.k_proj.input_scale", + "model.layers.0.self_attn.k_proj.input_scale", + ), + # regular weights untouched + ( + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + ), + ], + ) + def test_remap(self, name, expected): + assert self._map(name) == expected + + @pytest.mark.parametrize( + "name", + [ + "model.layers.0.self_attn.k_scale", + "model.layers.0.self_attn.k_proj.k_scale", + "model.layers.0.self_attn.qkv_proj.v_scale", + "model.layers.0.mixer.k_proj.k_scale", + ], + ) + def test_idempotent(self, name): + once = self._map(name) + assert once is not None + assert self._map(once) == once + + def test_composes_with_qkv_mapper(self): + """Applied together with a model's qkv/gate_up mapper, the regex scale + rules run before the substr rename, so scales are normalized to `.attn.` + and regular projections are still fused correctly.""" + from vllm.model_executor.models.utils import WeightsMapper + + model_mapper = WeightsMapper( + orig_to_new_substr={ + ".q_proj": ".qkv_proj.q", + ".k_proj": ".qkv_proj.k", + ".v_proj": ".qkv_proj.v", + } + ) + # AutoWeightsLoader does `mapper |= cache_scale_mapper` + combined = model_mapper | self._mapper() + + assert ( + combined._map_name("model.layers.0.self_attn.q_proj.weight") + == "model.layers.0.self_attn.qkv_proj.q.weight" + ) + assert ( + combined._map_name("model.layers.0.self_attn.k_proj.k_scale") + == "model.layers.0.self_attn.attn.k_scale" + ) + assert ( + combined._map_name("model.layers.0.self_attn.k_scale") + == "model.layers.0.self_attn.attn.k_scale" + ) + + if __name__ == "__main__": test_download_weights_from_hf() diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index c1aee79bec2..7082b7287d8 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -128,9 +128,13 @@ class WorkerLoRAManager: peft_helper.validate_legal(self.lora_config) # For some models like Qwen2VL, we need to use hf_to_vllm_mapper - # to ensure correct loading of lora weights. + # to ensure correct loading of lora weights. Drop the QKV/MLP fusion + # substr maps so constituent names (e.g. `q_proj`) survive for the + # LoRA manager to pack, while keeping genuine renames/prefixes. model = self._adapter_manager.model hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None) + if hf_to_vllm_mapper is not None: + hf_to_vllm_mapper = hf_to_vllm_mapper.get_unstacked_mapper() # Get model-defined prefixes to skip during LoRA loading. lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 48c1902e29a..e487b91e989 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -3,9 +3,12 @@ import itertools from abc import abstractmethod +from collections.abc import Iterable +from typing import Any import torch from torch.nn.parameter import Parameter +from typing_extensions import TypeIs import vllm.envs as envs from vllm.distributed import ( @@ -632,31 +635,31 @@ class MergedColumnParallelLinear(ColumnParallelLinear): disable_tp=disable_tp, ) - def validate_shard_id(self, loaded_shard_id: int | tuple[int, ...] | None): - if loaded_shard_id is None: - return - if isinstance(loaded_shard_id, tuple): - for idx in loaded_shard_id: + def validate_shard_id(self, shard_id: Any) -> TypeIs[int | tuple[int, ...] | None]: + if isinstance(shard_id, int): + if shard_id < 0 or shard_id >= len(self.output_sizes): + raise ValueError( + f"Shard id should be between 0 and {len(self.output_sizes) - 1}. " + f"Got shard id {shard_id}." + ) + return True + if shard_id is None: + return True + if isinstance(shard_id, tuple): + for idx in shard_id: if not (0 <= idx < len(self.output_sizes)): raise ValueError( f"Shard id index {idx} should be between 0 and " - f"{len(self.output_sizes) - 1}. Got shard id {loaded_shard_id}." + f"{len(self.output_sizes) - 1}. Got shard id {shard_id}." ) - if len(loaded_shard_id) > 1 and any( - b - a != 1 for a, b in zip(loaded_shard_id[:-1], loaded_shard_id[1:]) + if len(shard_id) > 1 and any( + b - a != 1 for a, b in zip(shard_id[:-1], shard_id[1:]) ): raise ValueError( "Shard id with multiple indices should be consecutive. " - f"Got shard id {loaded_shard_id}." + f"Got shard id {shard_id}." ) - return - elif isinstance(loaded_shard_id, int): - if loaded_shard_id < 0 or loaded_shard_id >= len(self.output_sizes): - raise ValueError( - f"Shard id should be between 0 and {len(self.output_sizes) - 1}. " - f"Got shard id {loaded_shard_id}." - ) - return + return True raise ValueError("This line should not be reached") def weight_loader( @@ -910,6 +913,31 @@ class MergedColumnParallelLinear(ColumnParallelLinear): tp_rank=self.tp_rank, ) + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[str]: + for name, loaded_weight in weights: + shard_id = getattr(loaded_weight, "shard_id", None) + self.validate_shard_id(shard_id) + # Load into self if name is not an attr of self or its submodules + param: Parameter + if "." in name: + submodule, _, attr = name.rpartition(".") + param = getattr(self.get_submodule(submodule), attr, self) + else: + param = getattr(self, name, self) + if param is None and name == "bias": + continue + param.weight_loader(param, loaded_weight, shard_id) + logger.debug( + "Loaded shard %s with shape %s into %s.%s", + shard_id, + loaded_weight.shape, + self.prefix, + name, + ) + yield name + class QKVParallelLinear(ColumnParallelLinear): """Linear layers for the attention's QKV transformation. @@ -996,17 +1024,13 @@ class QKVParallelLinear(ColumnParallelLinear): disable_tp=disable_tp, ) - def validate_shard_id(self, loaded_shard_id: str | None): - if loaded_shard_id is None: - return - if isinstance(loaded_shard_id, str): - if loaded_shard_id not in ["q", "k", "v"]: - raise ValueError( - "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', " - f"got shard id {loaded_shard_id}." - ) - return - raise ValueError("This line should not be reached") + def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]: + if shard_id in {"q", "k", "v"} or shard_id is None: + return True + raise ValueError( + "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', " + f"got shard id {shard_id}." + ) def _get_shard_offset_mapping(self, loaded_shard_id: str): shard_offset_mapping = { @@ -1302,6 +1326,31 @@ class QKVParallelLinear(ColumnParallelLinear): assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[str]: + for name, loaded_weight in weights: + shard_id = getattr(loaded_weight, "shard_id", None) + self.validate_shard_id(shard_id) + # Load into self if name is not an attr of self or its submodules + param: Parameter + if "." in name: + submodule, _, attr = name.rpartition(".") + param = getattr(self.get_submodule(submodule), attr, self) + else: + param = getattr(self, name, self) + if param is None and name == "bias": + continue + param.weight_loader(param, loaded_weight, shard_id) + logger.debug( + "Loaded shard %s with shape %s into %s.%s", + shard_id, + loaded_weight.shape, + self.prefix, + name, + ) + yield name + class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear): """QKV projection fused with a lightning-indexer's index_q/index_k. @@ -1387,15 +1436,14 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear): prefix=prefix, ) - def validate_shard_id(self, loaded_shard_id: str | None) -> None: - if loaded_shard_id is None: - return - if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"): - raise ValueError( - "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of " - "'q', 'k', 'v', 'index_q', 'index_k'; got " - f"{loaded_shard_id}." - ) + def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]: + if shard_id in {"q", "k", "v", "index_q", "index_k"} or shard_id is None: + return True + raise ValueError( + "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of " + "'q', 'k', 'v', 'index_q', 'index_k'; got " + f"{shard_id}." + ) def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None: h = self.head_size diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 9b18bdc132e..ad7aea175de 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -5,6 +5,7 @@ import inspect from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any +import regex as re import torch from torch import nn from transformers import PretrainedConfig @@ -19,10 +20,12 @@ else: class QuantizeMethodBase(ABC): """Base class for different quantized methods.""" - # Whether this method creates weights on meta device for online quantization. - # When True, weights are created on meta device and quantized layer-wise - # in process_weights_after_loading, reducing peak memory during loading. uses_meta_device: bool = False + """ + Whether this method creates weights on meta device for online quantization. + When True, weights are created on meta device and quantized layer-wise + in process_weights_after_loading, reducing peak memory during loading. + """ @abstractmethod def create_weights( @@ -84,6 +87,18 @@ def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> class QuantizationConfig(ABC): """Base class for quantization configs.""" + _ignore_unexpected_suffixes = ( + ".q_scale", + ".k_scale", + ".v_scale", + ".q_zero_point", + ".k_zero_point", + ".v_zero_point", + ) + """Suffixes of quantization parameters that may be present in the checkpoint but + not in the model, and should be ignored if unexpected during loading. These are used + after remapping, so should be in vLLM format (e.g. .q_scale, not .q.scale).""" + def __init__(self): super().__init__() # mapping is updated by models as they initialize @@ -176,14 +191,40 @@ class QuantizationConfig(ABC): """ raise NotImplementedError - def get_cache_scale_mapper(self) -> "WeightsMapper | None": + @staticmethod + def get_cache_scale_mapper() -> "WeightsMapper": """Mapping from checkpoint KV-cache scale names to vLLM scale names. Returning a mapper here causes `AutoWeightsLoader` to apply it to the weight stream automatically; individual model `load_weights` methods do not need to know about KV-cache scales. """ - return None + from vllm.model_executor.models.utils import WeightsMapper + + orig_to_new_regex = { + # Deprecated fused kv_scale -> attn.k_scale + re.compile(r"\.kv_scale$"): r".attn.k_scale", + # ModelOpt: .self_attn.{k,v}_proj.{k,v}_scale -> .self_attn.attn.* + re.compile(r"\.self_attn\.[kv]_proj\.([kv])_scale$"): ( + r".self_attn.attn.\1_scale" + ), + # Fused QKV / qkqkv proj: .self_attn.qk(qk)v_proj.{k,v}_scale -> attn + re.compile(r"\.self_attn\.qk(?:qk)?v_proj\.([kv])_scale$"): ( + r".self_attn.attn.\1_scale" + ), + # NemotronH: .mixer.{k,v}_proj.{k,v}_scale -> .mixer.attn.* + re.compile(r"\.mixer\.[kv]_proj\.([kv])_scale$"): r".mixer.attn.\1_scale", + # HYV3: .self_attn.q.scale -> .self_attn.attn.q_scale + re.compile(r"\.self_attn\.q\.scale$"): r".self_attn.attn.q_scale", + # HYV3: .self_attn.{k,v}_cache.scale -> .self_attn.attn.{k,v}_scale + re.compile(r"\.self_attn\.([kv])_cache\.scale$"): ( + r".self_attn.attn.\1_scale" + ), + # Default: .{q,k,v}_scale -> .attn.{q,k,v}_scale (unless already .attn) + re.compile(r"(? "WeightsMapper": + @staticmethod + def get_cache_scale_mapper() -> "WeightsMapper": """Map compressed-tensors KV-cache scale names to vLLM names.""" from vllm.model_executor.models.utils import WeightsMapper - return WeightsMapper( - orig_to_new_suffix={ - ".k_proj.output_scale": ".attn.k_scale", - ".v_proj.output_scale": ".attn.v_scale", - ".q_proj.output_scale": ".attn.q_scale", - ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", - } - ) + orig_to_new_suffix = { + ".k_proj.output_scale": ".attn.k_scale", + ".v_proj.output_scale": ".attn.v_scale", + ".q_proj.output_scale": ".attn.q_scale", + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", + } + cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix) + return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper() class CopyNumelCounter(TorchDispatchMode): diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index 9051214cf9d..fbd61e28cd2 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -679,16 +679,17 @@ class QuarkConfig(QuantizationConfig): return scheme - def get_cache_scale_mapper(self) -> "WeightsMapper": + @staticmethod + def get_cache_scale_mapper() -> "WeightsMapper": """Map Quark KV-cache scale names to vLLM names.""" - return WeightsMapper( - orig_to_new_suffix={ - ".k_proj.output_scale": ".attn.k_scale", - ".v_proj.output_scale": ".attn.v_scale", - ".q_proj.output_scale": ".attn.q_scale", - ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", - } - ) + orig_to_new_suffix = { + ".k_proj.output_scale": ".attn.k_scale", + ".v_proj.output_scale": ".attn.v_scale", + ".q_proj.output_scale": ".attn.q_scale", + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", + } + cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix) + return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper() class QuarkLinearMethod(LinearMethodBase): diff --git a/vllm/model_executor/model_loader/bitsandbytes_loader.py b/vllm/model_executor/model_loader/bitsandbytes_loader.py index 064a74023a2..55b5d617a73 100644 --- a/vllm/model_executor/model_loader/bitsandbytes_loader.py +++ b/vllm/model_executor/model_loader/bitsandbytes_loader.py @@ -576,7 +576,8 @@ class BitsAndBytesModelLoader(BaseModelLoader): # For some models like Molmo, we need to use hf_to_vllm_mapper # to ensure correct loading of weights. if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None): - self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name) + unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper() + self.weight_mapper = lambda name, m=unstacked_mapper: m._map_name(name) self._get_bnb_target_modules(model) self._classify_module_sharding(model) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 6cf1c19cba4..d0d26fed3e6 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -131,8 +131,11 @@ def initialize_online_processing(layer: torch.nn.Module): # Track loading progress to determine when to process/copy info.load_numel = 0 info.load_numel_total = get_layer_size(layer) + _wrap_parameters_weight_loader(layer) - # Wrap each parameter's weight loader + +def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None: + """Wrap each parameter's weight loader.""" # Note that nested wrapping will occur for shared tensors for name, tensor in get_layer_tensors(layer).items(): if name in SKIP_TENSORS: @@ -168,6 +171,12 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla logger.debug("%s: Excessive loading", layer.__class__.__name__) return + # Re-run on each load: layers may register parameters later (e.g., `bias`). + # Wrap late parameters and refresh `load_numel_total` so processing waits + # until all parameters are loaded. + info.load_numel_total = get_layer_size(layer) + _wrap_parameters_weight_loader(layer) + # Bind and normalize arguments bound_args = loader_signature.bind(*args, **kwargs) bound_args.apply_defaults() diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index fc279c7e9c7..fc59acf3d35 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -290,6 +290,6 @@ def configure_quant_config( # pass mappings by reference to quant_config if hf_to_vllm_mapper is not None: - quant_config.apply_vllm_mapper(hf_to_vllm_mapper) + quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_unstacked_mapper()) if packed_mapping is not None: quant_config.packed_modules_mapping = packed_mapping diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index d25c954fc19..c32a903bba8 100644 --- a/vllm/model_executor/models/arcee.py +++ b/vllm/model_executor/models/arcee.py @@ -26,10 +26,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import ( @@ -42,7 +38,7 @@ from .interfaces import ( from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -276,67 +272,6 @@ class ArceeModel(nn.Module, EagleModelMixin): return hidden_states, aux_hidden_states return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights, mapping q/k/v projections to fused qkv_proj.""" - stacked_params_mapping = [ - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - ] - - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - continue - - if "scale" in name or "zero_point" in name: - remapped_name = maybe_remap_kv_scale_name(name, params_dict) - if remapped_name is None: - continue - name = remapped_name - - mapped = False - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - - name = name.replace(weight_name, param_name) - - if name.endswith(".bias") and name not in params_dict: - mapped = True - break - - if is_pp_missing_parameter(name, self): - mapped = True - break - - param = params_dict[name] - weight_loader = param.weight_loader # type: ignore[attr-defined] - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - mapped = True - break - - if mapped: - continue - - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - return loaded_params - class ArceeForCausalLM( nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 @@ -344,6 +279,14 @@ class ArceeForCausalLM( """Arcee Model for causal language modeling, integrated with vLLM runtime.""" + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + } + ) # Map fused module names to their submodule components # (for quantization and LoRA) packed_modules_mapping = { @@ -420,4 +363,4 @@ class ArceeForCausalLM( ) # AutoWeightLoader handles weight name remapping, including fusing # separate q_proj, k_proj, v_proj into qkv_proj - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/chatglm.py b/vllm/model_executor/models/chatglm.py index c5d857e7c3d..4363188ff6e 100644 --- a/vllm/model_executor/models/chatglm.py +++ b/vllm/model_executor/models/chatglm.py @@ -30,7 +30,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.chatglm import ChatGLMConfig @@ -38,7 +37,6 @@ from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant from .utils import ( AutoWeightsLoader, WeightsMapper, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -316,12 +314,9 @@ class GLMTransformer(nn.Module): @support_torch_compile class ChatGLMModel(nn.Module, SupportsQuant): - packed_modules_mapping = { - "linear_proj.merged_proj": [ - "linear_proj.gate_proj", - "linear_proj.dense_h_to_4h", - ] - } + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={".word_embeddings": ""}, + ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -386,47 +381,11 @@ class ChatGLMModel(nn.Module, SupportsQuant): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("linear_proj.merged_proj", "linear_proj.gate_proj", 0), - ("linear_proj.merged_proj", "linear_proj.dense_h_to_4h", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if "rotary_pos_emb.inv_freq" in name: - continue - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class ChatGLMBaseModel(nn.Module): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={".word_embeddings": ""}, - ) - def __init__( self, *, @@ -467,7 +426,7 @@ class ChatGLMBaseModel(nn.Module): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loader.load_weights(weights) class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP, SupportsQuant): diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py index 7b57c739ffe..64ec0d6dd54 100644 --- a/vllm/model_executor/models/cohere_eagle.py +++ b/vllm/model_executor/models/cohere_eagle.py @@ -14,7 +14,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.commandr import ( CohereDecoderLayer, CohereForCausalLM, @@ -134,42 +133,6 @@ class CohereEagleModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states, hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if name.endswith(".bias") and name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if name.endswith(".bias") and name not in params_dict: - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class EagleCohereForCausalLM(CohereForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -225,7 +188,9 @@ class EagleCohereForCausalLM(CohereForCausalLM): ), ) - loaded_weight_names = loader.load_weights(map(_track_and_forward, weights)) + loaded_weight_names = loader.load_weights( + map(_track_and_forward, weights), mapper=self.hf_to_vllm_mapper + ) # Embed tokens are tied with the target model and therefore not # present in the EAGLE checkpoint; mark them as loaded explicitly to diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 66adb9a3ca7..3d5120b4d07 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -45,8 +45,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, row_parallel_weight_loader, ) from vllm.model_executor.utils import set_weight_attrs @@ -58,7 +56,6 @@ from .utils import ( AutoWeightsLoader, WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -341,60 +338,21 @@ class CohereModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, shard_name, shard_id in stacked_params_mapping: - if shard_name not in name: - continue - name = name.replace(shard_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } # LoRA specific attributes embedding_modules = {"embed_tokens": "input_embeddings"} diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py index 7796c3da331..79314a7b931 100644 --- a/vllm/model_executor/models/exaone.py +++ b/vllm/model_executor/models/exaone.py @@ -50,17 +50,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -370,70 +366,21 @@ class ExaoneModel(nn.Module): hidden_states, _ = self.ln_f(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".c_fc_0", 0), - (".gate_up_proj", ".c_fc_1", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".c_fc_0": (".gate_up_proj", 0), + ".c_fc_1": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "c_fc_0", - "c_fc_1", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["c_fc_0", "c_fc_1"], } # LoRA specific attributes @@ -506,4 +453,4 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP): # processed with quantization, LoRA, fine-tuning, etc. skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index cc1dcf197f7..dc88c15fc01 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -46,10 +46,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.config import set_default_rope_theta @@ -57,8 +53,8 @@ from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -368,70 +364,21 @@ class Exaone4Model(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } # LoRA specific attributes @@ -503,4 +450,4 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): # processed with quantization, LoRA, fine-tuning, etc. skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/fairseq2_llama.py b/vllm/model_executor/models/fairseq2_llama.py index ca0e7e64df5..e898034fbfa 100644 --- a/vllm/model_executor/models/fairseq2_llama.py +++ b/vllm/model_executor/models/fairseq2_llama.py @@ -79,10 +79,8 @@ class Fairseq2LlamaForCausalLM(LlamaForCausalLM): skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) return loader.load_weights( - ( - self.reshape_fairseq2_weights(name, loaded_weight, params) - for name, loaded_weight in weights - ) + self.reshape_fairseq2_weights(name, loaded_weight, params) + for name, loaded_weight in weights ) def flag_sharded_weights(self, params: dict[str, Parameter]): diff --git a/vllm/model_executor/models/gemma.py b/vllm/model_executor/models/gemma.py index 6e35020a6ea..949799fa654 100644 --- a/vllm/model_executor/models/gemma.py +++ b/vllm/model_executor/models/gemma.py @@ -42,13 +42,12 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors -from .interfaces import SupportsLoRA, SupportsPP +from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -324,56 +323,21 @@ class GemmaModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, shard_name, shard_id in stacked_params_mapping: - if shard_name not in name: - continue - name = name.replace(shard_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP): +class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -421,4 +385,4 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma2.py b/vllm/model_executor/models/gemma2.py index 733eb3ed3c1..da5161ffa01 100644 --- a/vllm/model_executor/models/gemma2.py +++ b/vllm/model_executor/models/gemma2.py @@ -39,17 +39,13 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -316,60 +312,21 @@ class Gemma2Model(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, shard_name, shard_id in stacked_params_mapping: - if shard_name not in name: - continue - name = name.replace(shard_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - return loaded_params - class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -418,4 +375,4 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py index 308c9c8a8ea..717bc62439a 100644 --- a/vllm/model_executor/models/gemma3.py +++ b/vllm/model_executor/models/gemma3.py @@ -44,18 +44,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -365,65 +361,18 @@ class Gemma3Model(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - # Check if this is a scale parameter that needs remapping first - if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): - # Try to remap the scale name first - remapped_name = maybe_remap_kv_scale_name(name, params_dict) - if remapped_name is not None and remapped_name in params_dict: - # Successfully remapped, use the remapped name - param = params_dict[remapped_name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(remapped_name) - continue - # If remapping failed, continue with normal processing - - for param_name, shard_name, shard_id in stacked_params_mapping: - if shard_name not in name: - continue - name = name.replace(shard_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - return loaded_params - class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -491,4 +440,4 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/glm4.py b/vllm/model_executor/models/glm4.py index 4587a692766..3a25f90ad2a 100644 --- a/vllm/model_executor/models/glm4.py +++ b/vllm/model_executor/models/glm4.py @@ -39,10 +39,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType @@ -52,7 +48,6 @@ from .llama import LlamaModel from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, maybe_prefix, ) @@ -237,73 +232,11 @@ class Glm4Model(LlamaModel): vllm_config=vllm_config, prefix=prefix, layer_type=Glm4DecoderLayer ) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) - if spec_layer is not None: - continue - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if "scale" in name or "zero_point" in name: - # Remapping the name of FP8 kv-scale or zero point. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -360,10 +293,15 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else [] + # Skip the speculative (MTP) layers, which are loaded by the + # draft model instead. + num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0) + skip_prefixes += [ + f"model.layers.{self.config.num_hidden_layers + i}." + for i in range(num_nextn_layers) + ] + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) return loader.load_weights(weights) diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py index 9d08df4df8d..2e3a301579d 100644 --- a/vllm/model_executor/models/glm4v.py +++ b/vllm/model_executor/models/glm4v.py @@ -61,6 +61,7 @@ from .interfaces import ( SupportsMultiModal, SupportsPP, ) +from .utils import WeightsMapper class GLMVImagePixelInputs(TensorSchema): @@ -376,6 +377,15 @@ class EVA2CLIPModel(nn.Module): class GLM4VModel(ChatGLMModel): + hf_to_vllm_mapper = ChatGLMModel.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + # Vision GLU projections + "linear_proj.gate_proj": ("linear_proj.merged_proj", 0), + "linear_proj.dense_h_to_4h": ("linear_proj.merged_proj", 1), + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) @@ -507,6 +517,9 @@ class GLM4VMultiModalProcessor(BaseMultiModalProcessor[GLM4VProcessingInfo]): class GLM4VForCausalLM( ChatGLMBaseModel, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE ): + # NOTE: we must bring this to the surface because GLM4VModel.hf_to_vllm_mapper + # contains non-stacking related mappings which LoRA/BnB needs to know about + hf_to_vllm_mapper = GLM4VModel.hf_to_vllm_mapper packed_modules_mapping = { "query_key_value": ["query_key_value"], "dense_h_to_4h": ["dense_h_to_4h"], diff --git a/vllm/model_executor/models/gpt_j.py b/vllm/model_executor/models/gpt_j.py index 30da9b4dea2..44dec873457 100644 --- a/vllm/model_executor/models/gpt_j.py +++ b/vllm/model_executor/models/gpt_j.py @@ -43,16 +43,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -239,51 +235,17 @@ class GPTJModel(nn.Module): hidden_states = self.ln_f(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "attn.bias" in name or "attn.masked_bias" in name: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class GPTJForCausalLM(nn.Module, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -329,5 +291,5 @@ class GPTJForCausalLM(nn.Module, SupportsPP): return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self, skip_substrs=["attn.bias", "attn.masked_bias"]) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index 7470e7e7381..c46fefbf889 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -49,17 +49,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors -from .interfaces import SupportsLoRA, SupportsPP +from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_layers, maybe_prefix, ) @@ -252,6 +248,17 @@ class GraniteDecoderLayer(nn.Module): @support_torch_compile class GraniteModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -322,66 +329,17 @@ class GraniteModel(nn.Module): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) -class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP): - packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], - } - +class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): + hf_to_vllm_mapper = GraniteModel.hf_to_vllm_mapper # LoRA specific attributes + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } embedding_modules = { "embed_tokens": "input_embeddings", "lm_head": "output_embeddings", diff --git a/vllm/model_executor/models/hyperclovax.py b/vllm/model_executor/models/hyperclovax.py index 2f54f78e758..8ba07926259 100644 --- a/vllm/model_executor/models/hyperclovax.py +++ b/vllm/model_executor/models/hyperclovax.py @@ -50,10 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.hyperclovax import HyperCLOVAXConfig @@ -61,7 +57,7 @@ from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -377,71 +373,21 @@ class HyperCLOVAXModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if "scale" in name or "zero_point" in name: - # Remapping the name of FP8 kv-scale or zero point. - remapped_name = maybe_remap_kv_scale_name(name, params_dict) - if remapped_name is None: - continue - name = remapped_name - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader # type: ignore[attr-defined] - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class HyperCLOVAXForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } # LoRA specific attributes @@ -536,4 +482,4 @@ class HyperCLOVAXForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=["lm_head."] if self.config.tie_word_embeddings else None, ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 29603318c15..f1d6d563738 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1033,7 +1033,8 @@ class SupportsQuant: if self.quant_config is None: return if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None: - self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper) + unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper() + self.quant_config.apply_vllm_mapper(unstacked_mapper) if self.packed_modules_mapping is not None: self.quant_config.packed_modules_mapping.update(self.packed_modules_mapping) diff --git a/vllm/model_executor/models/internlm2.py b/vllm/model_executor/models/internlm2.py index 6b1712ede32..81487f9cad5 100644 --- a/vllm/model_executor/models/internlm2.py +++ b/vllm/model_executor/models/internlm2.py @@ -35,15 +35,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors -from .interfaces import SupportsLoRA, SupportsPP +from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant from .interfaces_base import default_pooling_type from .utils import ( AutoWeightsLoader, StageMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -249,6 +248,14 @@ class InternLMDecoderLayer(nn.Module): @support_torch_compile class InternLM2Model(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".w1": (".gate_up_proj", 0), + ".w3": (".gate_up_proj", 1), + } + ) + def __init__( self, *, @@ -310,43 +317,12 @@ class InternLM2Model(nn.Module): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("gate_up_proj", "w1", 0), - ("gate_up_proj", "w3", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) -class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): +class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsQuant): + hf_to_vllm_mapper = InternLM2Model.hf_to_vllm_mapper packed_modules_mapping = { "wqkv": ["wqkv"], "gate_up_proj": ["w1", "w3"], diff --git a/vllm/model_executor/models/jais2.py b/vllm/model_executor/models/jais2.py index 325d5249289..95b8c3ee44f 100644 --- a/vllm/model_executor/models/jais2.py +++ b/vllm/model_executor/models/jais2.py @@ -51,18 +51,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -366,61 +362,16 @@ class Jais2Model(nn.Module): hidden_states, _ = self.norm(hidden_states + residual), residual return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if "scale" in name: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + } + ) packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], } @@ -490,4 +441,4 @@ class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 2b07937df08..82a53440402 100644 --- a/vllm/model_executor/models/jina.py +++ b/vllm/model_executor/models/jina.py @@ -254,5 +254,6 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling): 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} + loader = AutoWeightsLoader(self.model, ignore_unexpected_prefixes=["lm_head."]) + weights = _merge_weights(weights) + return loader.load_weights(weights, mapper=self.model.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index a54801e6458..bb223a31146 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -52,10 +52,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType @@ -67,12 +63,13 @@ from .interfaces import ( SupportsEagle3, SupportsLoRA, SupportsPP, + SupportsQuant, ) from .utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -345,6 +342,17 @@ class LlamaDecoderLayer(nn.Module): }, ) class LlamaModel(nn.Module, EagleModelMixin): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + def __init__( self, *, @@ -431,67 +439,25 @@ class LlamaModel(nn.Module, EagleModelMixin): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if "scale" in name or "zero_point" in name: - # Remapping the name of FP8 kv-scale or zero point. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class LlamaForCausalLM( - LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 + LocalArgmaxMixin, + nn.Module, + SupportsLoRA, + SupportsPP, + SupportsEagle, + SupportsEagle3, + SupportsQuant, ): + hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper + # LoRA specific attributes packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], } - - # LoRA specific attributes embedding_modules = { "embed_tokens": "input_embeddings", "lm_head": "output_embeddings", diff --git a/vllm/model_executor/models/mamba.py b/vllm/model_executor/models/mamba.py index ec2a7255eb6..6a77a58abf4 100644 --- a/vllm/model_executor/models/mamba.py +++ b/vllm/model_executor/models/mamba.py @@ -26,7 +26,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import ( HasInnerState, IsAttentionFree, @@ -37,7 +36,7 @@ from vllm.sequence import IntermediateTensors from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -170,28 +169,12 @@ class MambaModel(nn.Module): return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "A_log" in name: - name = name.replace("A_log", "A") - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class MambaForCausalLM( nn.Module, HasInnerState, IsAttentionFree, SupportsPP, SupportsMambaPrefixCaching ): + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".A_log": ".A"}) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config @@ -279,4 +262,4 @@ class MambaForCausalLM( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/mamba2.py b/vllm/model_executor/models/mamba2.py index deb20852a26..343111ee015 100644 --- a/vllm/model_executor/models/mamba2.py +++ b/vllm/model_executor/models/mamba2.py @@ -25,7 +25,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import ( HasInnerState, IsAttentionFree, @@ -35,7 +34,7 @@ from vllm.sequence import IntermediateTensors from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -167,29 +166,12 @@ class Mamba2Model(nn.Module): return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "A_log" in name: - name = name.replace("A_log", "A") - - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class Mamba2ForCausalLM( nn.Module, HasInnerState, IsAttentionFree, SupportsMambaPrefixCaching ): + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".A_log": ".A"}) + @classmethod def get_mamba_state_dtype_from_config( cls, @@ -292,4 +274,4 @@ class Mamba2ForCausalLM( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/mimo.py b/vllm/model_executor/models/mimo.py index 4f67d468ace..e4247fa8d8d 100644 --- a/vllm/model_executor/models/mimo.py +++ b/vllm/model_executor/models/mimo.py @@ -38,14 +38,10 @@ from vllm.distributed import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2Model from vllm.sequence import IntermediateTensors -from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix +from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix logger = init_logger(__name__) @@ -89,50 +85,6 @@ class MiMoModel(Qwen2Model): hidden_states = hidden_states + residual return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - ("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), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "mtp_layers" in name: - continue - if "rotary_emb.inv_freq" in name: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -167,6 +119,13 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module): self.model.make_empty_intermediate_tensors ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else [] + # MTP layers are loaded by the draft model, not the main model. + skip_prefixes.append("model.mtp_layers.") + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + return loader.load_weights(weights) + def compute_logits( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/models/mistral_eagle.py b/vllm/model_executor/models/mistral_eagle.py index 8865742d649..75d1ebb91a8 100644 --- a/vllm/model_executor/models/mistral_eagle.py +++ b/vllm/model_executor/models/mistral_eagle.py @@ -108,11 +108,6 @@ class EagleMistralModel(MistralModel): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states, hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Pretend embed_tokens is loaded; the actual weight is shared - # from the target model at runtime by `load_eagle_model`. - return super().load_weights(weights) | {"embed_tokens.weight"} - class EagleMistralForCausalLM(MistralForCausalLM): mistral_mapping = MistralForCausalLM.mistral_mapping | { @@ -166,3 +161,8 @@ class EagleMistralForCausalLM(MistralForCausalLM): multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Pretend embed_tokens is loaded; the actual weight is shared + # from the target model at runtime by `load_eagle_model`. + return super().load_weights(weights) | {"model.embed_tokens.weight"} diff --git a/vllm/model_executor/models/mpt.py b/vllm/model_executor/models/mpt.py index 85933626cd3..8e509fbcb4c 100644 --- a/vllm/model_executor/models/mpt.py +++ b/vllm/model_executor/models/mpt.py @@ -27,13 +27,11 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .interfaces import SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -274,21 +272,6 @@ class MPTModel(nn.Module): hidden_states = self.norm_f(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class MPTForCausalLM(nn.Module, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py index e276d5368ad..6f0b61205b3 100644 --- a/vllm/model_executor/models/nemotron.py +++ b/vllm/model_executor/models/nemotron.py @@ -47,10 +47,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.nemotron import NemotronConfig @@ -58,7 +54,7 @@ from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -366,58 +362,18 @@ class NemotronModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], } # LoRA specific attributes @@ -485,4 +441,4 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/nemotron_nas.py b/vllm/model_executor/models/nemotron_nas.py index 06a2096ec69..5a5f0e77739 100644 --- a/vllm/model_executor/models/nemotron_nas.py +++ b/vllm/model_executor/models/nemotron_nas.py @@ -42,10 +42,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.model_executor.models.llama import LlamaAttention, LlamaMLP from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType @@ -54,7 +50,7 @@ from .interfaces import HasNoOps, SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -315,60 +311,18 @@ class DeciModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if "scale" in name or "zero_point" in name: - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -462,4 +416,4 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/olmo.py b/vllm/model_executor/models/olmo.py index 541f60c2c40..e62bd39238b 100644 --- a/vllm/model_executor/models/olmo.py +++ b/vllm/model_executor/models/olmo.py @@ -48,13 +48,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -301,59 +300,25 @@ class OlmoModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA): """ Extremely barebones HF model wrapper. """ + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -410,4 +375,4 @@ class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA): ["lm_head.weight"] if self.config.tie_word_embeddings else None ), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/olmo2.py b/vllm/model_executor/models/olmo2.py index ad04b258bde..489ec2616cb 100644 --- a/vllm/model_executor/models/olmo2.py +++ b/vllm/model_executor/models/olmo2.py @@ -52,12 +52,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP from vllm.model_executor.models.utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -343,58 +342,25 @@ class Olmo2Model(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if is_pp_missing_parameter(name, self): - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader # type: ignore - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): """ Extremely barebones HF model wrapper. """ + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -451,4 +417,4 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): ["lm_head.weight"] if self.config.tie_word_embeddings else None ), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py index 81653b9516a..32bb532f5c5 100644 --- a/vllm/model_executor/models/opt.py +++ b/vllm/model_executor/models/opt.py @@ -44,14 +44,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, WeightsMapper, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -325,53 +323,23 @@ class OPTModel(nn.Module): input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds ) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + }, + orig_to_new_prefix={ + "decoder.": "model.decoder.", + }, + ) packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], } - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - "decoder.": "model.decoder.", - } - ) - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/orion.py b/vllm/model_executor/models/orion.py index 3cacb9d61cd..0871c347ac5 100644 --- a/vllm/model_executor/models/orion.py +++ b/vllm/model_executor/models/orion.py @@ -32,13 +32,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .interfaces import SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -277,45 +276,19 @@ class OrionModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class OrionForCausalLM(nn.Module, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -362,4 +335,4 @@ class OrionForCausalLM(nn.Module, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/ouro.py b/vllm/model_executor/models/ouro.py index 503d4b5c834..527eeaa13bc 100644 --- a/vllm/model_executor/models/ouro.py +++ b/vllm/model_executor/models/ouro.py @@ -51,16 +51,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType from .interfaces import SupportsLoRA from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, make_empty_intermediate_tensors_factory, make_layers, @@ -376,65 +373,21 @@ class OuroModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if name.endswith("scale"): - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - if weight_loader == default_weight_loader: - weight_loader(param, loaded_weight) - else: - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class OuroForCausalLM(nn.Module, SupportsLoRA): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -492,4 +445,4 @@ class OuroForCausalLM(nn.Module, SupportsLoRA): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/phi.py b/vllm/model_executor/models/phi.py index 75c42c0d393..61c243aadf2 100644 --- a/vllm/model_executor/models/phi.py +++ b/vllm/model_executor/models/phi.py @@ -62,13 +62,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -257,55 +256,18 @@ class PhiModel(nn.Module): return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # pylint: disable=E1136 - - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ] + "qkv_proj": ["q_proj", "k_proj", "v_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -360,4 +322,4 @@ class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index 9c39c649708..182b9758308 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/vllm/model_executor/models/qwen2.py @@ -54,10 +54,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta from vllm.v1.attention.backend import AttentionType @@ -68,12 +64,13 @@ from .interfaces import ( SupportsEagle3, SupportsLoRA, SupportsPP, + SupportsQuant, ) from .utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -323,6 +320,17 @@ class Qwen2DecoderLayer(nn.Module): } ) class Qwen2Model(nn.Module, EagleModelMixin): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + def __init__( self, *, @@ -426,72 +434,17 @@ class Qwen2Model(nn.Module, EagleModelMixin): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - if name.endswith("scale"): - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - if weight_loader == default_weight_loader: - weight_loader(param, loaded_weight) - else: - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class Qwen2ForCausalLM( - nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 + nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3, SupportsQuant ): + hf_to_vllm_mapper = Qwen2Model.hf_to_vllm_mapper packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm/model_executor/models/qwen2_rm.py b/vllm/model_executor/models/qwen2_rm.py index cdf1a327efe..f2603431668 100644 --- a/vllm/model_executor/models/qwen2_rm.py +++ b/vllm/model_executor/models/qwen2_rm.py @@ -28,16 +28,10 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP): is_pooling_model = True pooler: Pooler + hf_to_vllm_mapper = Qwen2Model.hf_to_vllm_mapper packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index b070eac3255..a21f5b3b89c 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -267,18 +267,11 @@ class Qwen3Model(Qwen2Model): class Qwen3ForCausalLM( LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 ): + hf_to_vllm_mapper = Qwen3Model.hf_to_vllm_mapper packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } - embedding_modules = { "embed_tokens": "input_embeddings", "lm_head": "output_embeddings", diff --git a/vllm/model_executor/models/rnj1.py b/vllm/model_executor/models/rnj1.py index 68c3722e2bc..2bcd2791981 100644 --- a/vllm/model_executor/models/rnj1.py +++ b/vllm/model_executor/models/rnj1.py @@ -30,18 +30,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, - is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -331,75 +327,21 @@ class Rnj1Model(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if ( - self.quant_config - and self.quant_config.get_name() == "gguf" - and name.endswith("norm.weight") - ): - loaded_weight -= 1 - - if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): - remapped_name = maybe_remap_kv_scale_name(name, params_dict) - if remapped_name is not None and remapped_name in params_dict: - param = params_dict[remapped_name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(remapped_name) - continue - - for param_name, shard_name, shard_id in stacked_params_mapping: - if shard_name not in name: - continue - name = name.replace(shard_name, param_name) - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if name.endswith(".bias") and name not in params_dict: - continue - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - return loaded_params - class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -457,4 +399,4 @@ class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py index 48147f7334e..d2c767846d7 100644 --- a/vllm/model_executor/models/seed_oss.py +++ b/vllm/model_executor/models/seed_oss.py @@ -49,10 +49,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.config import set_default_rope_theta from vllm.v1.attention.backend import AttentionType @@ -61,7 +57,7 @@ from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -362,61 +358,21 @@ class SeedOssModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -477,4 +433,4 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py index fcb2ae429cb..478a61da675 100644 --- a/vllm/model_executor/models/solar.py +++ b/vllm/model_executor/models/solar.py @@ -48,17 +48,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -347,66 +343,22 @@ class SolarModel(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } - # LoRA specific attributes embedding_modules = { "embed_tokens": "input_embeddings", @@ -468,4 +420,4 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/stablelm.py b/vllm/model_executor/models/stablelm.py index 034c9c18ff7..58758b11cdd 100644 --- a/vllm/model_executor/models/stablelm.py +++ b/vllm/model_executor/models/stablelm.py @@ -45,13 +45,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .interfaces import SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -266,45 +265,19 @@ class StableLMEpochModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class StablelmForCausalLM(nn.Module, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -351,4 +324,4 @@ class StablelmForCausalLM(nn.Module, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/starcoder2.py b/vllm/model_executor/models/starcoder2.py index 5f08a59e236..08463011fe0 100644 --- a/vllm/model_executor/models/starcoder2.py +++ b/vllm/model_executor/models/starcoder2.py @@ -45,16 +45,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.sequence import IntermediateTensors from .interfaces import SupportsPP from .utils import ( AutoWeightsLoader, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -272,41 +268,17 @@ class Starcoder2Model(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ] - - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class Starcoder2ForCausalLM(nn.Module, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -362,4 +334,4 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP): ["lm_head.weight"] if self.config.tie_word_embeddings else None ), ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/step1.py b/vllm/model_executor/models/step1.py index 07653fa6b37..c18bf8a3c35 100644 --- a/vllm/model_executor/models/step1.py +++ b/vllm/model_executor/models/step1.py @@ -30,7 +30,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import ( EagleModelMixin, SupportsEagle, @@ -40,7 +39,7 @@ from vllm.model_executor.models.interfaces import ( from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, - is_pp_missing_parameter, + WeightsMapper, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, @@ -48,11 +47,6 @@ from vllm.model_executor.models.utils import ( from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType -STEP_PACKED_MODULES_MAPPING = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"], -} - def _get_step_alibi_slopes(total_num_heads: int) -> torch.Tensor: """Reference ALiBi slopes used by Step models.""" @@ -242,42 +236,6 @@ class StepDecoderLayer(nn.Module): hidden_states = self.mlp(hidden_states) return hidden_states, residual - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - (".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), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) # type: ignore[name-defined] - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class StepDecoderModel(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -354,7 +312,20 @@ class StepDecoderModel(nn.Module, EagleModelMixin): class Step1ForCausalLM(nn.Module, SupportsPP, SupportsEagle, SupportsEagle3): - packed_modules_mapping = STEP_PACKED_MODULES_MAPPING + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -413,4 +384,4 @@ class Step1ForCausalLM(nn.Module, SupportsPP, SupportsEagle, SupportsEagle3): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 55d94600497..4402d180ca0 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -158,9 +158,6 @@ class Base( "Transformers modeling backend does " "not support MXFP4 quantization yet." ) - # Skip loading extra bias for GPTQ models. - if "gptq" in quant_method_name: - self.ignore_unexpected_suffixes.append(".bias") self._patch_config() from_config_kwargs = dict( diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 730dc81ed21..6f4524400c0 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -4,8 +4,8 @@ import itertools from collections.abc import Callable, Iterable, Mapping from contextlib import contextmanager -from dataclasses import dataclass, field -from typing import Any, Literal, Protocol, overload +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, overload import regex as re import torch @@ -19,9 +19,6 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.base_config import ( - QuantizationConfig, -) from vllm.model_executor.model_loader.reload import ( support_quantized_model_reload_from_hp_weights, ) @@ -35,8 +32,13 @@ from vllm.utils.torch_utils import ( direct_register_custom_op, ) +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization import QuantizationConfig + logger = init_logger(__name__) +ShardId: TypeAlias = str | int | tuple[int, ...] + @dataclass class WeightsMapper: @@ -47,6 +49,7 @@ class WeightsMapper: orig_to_new_renamings: list[Any] = field(default_factory=list) orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict) + orig_to_new_stacked: Mapping[str, tuple[str, ShardId]] = field(default_factory=dict) orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict) orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict) @@ -59,11 +62,36 @@ class WeightsMapper: ], orig_to_new_regex={**self.orig_to_new_regex, **other.orig_to_new_regex}, orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr}, + orig_to_new_stacked={ + **self.orig_to_new_stacked, + **other.orig_to_new_stacked, + }, orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix}, orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix}, ) def _map_name(self, key: str) -> str | None: + """Map a weight name (backward-compatible wrapper that discards shard_id).""" + result = self._map_name_with_shard(key) + return result[0] if result is not None else None + + def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None: + """Map a weight name and extract any shard_id metadata. + + Returns: + (mapped_name, shard_id) if the name should be kept. + None if the name should be dropped. + """ + # Deprecation warnings + if key.endswith(".kv_scale"): + logger.warning_once( + "DEPRECATED. Found kv_scale in the checkpoint. " + "This format is deprecated in favor of separate k_scale and " + "v_scale tensors and will be removed in a future release. " + "Functionally, we will remap kv_scale to k_scale and duplicate " + "k_scale to v_scale" + ) + for renaming in self.orig_to_new_renamings: key, _ = renaming.rename_source_key(key) @@ -81,6 +109,12 @@ class WeightsMapper: key = key.replace(substr, new_key, 1) + shard_id: ShardId | None = None + for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items(): + if substr in key: + key = key.replace(substr, new_key, 1) + shard_id = new_shard_id + for prefix, new_key in self.orig_to_new_prefix.items(): if key.startswith(prefix): if new_key is None: @@ -95,16 +129,19 @@ class WeightsMapper: key = new_key.join(key.rsplit(suffix, 1)) - return key + return key, shard_id def apply( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[tuple[str, torch.Tensor]]: - return ( - (out_name, data) - for name, data in weights - if (out_name := self._map_name(name)) is not None - ) + for name, data in weights: + result = self._map_name_with_shard(name) + if result is None: + continue + out_name, shard_id = result + if shard_id is not None: + data.shard_id = shard_id + yield out_name, data def apply_list(self, values: list[str]) -> list[str]: return [ @@ -120,6 +157,15 @@ class WeightsMapper: if (out_name := self._map_name(name)) is not None } + def get_unstacked_mapper(self) -> "WeightsMapper": + """Mapper variant that drops stacked maps, keeping all genuine renames/prefixes. + + Consumers that reference the checkpoint's *unstacked* module names (LoRA name + parsing and the quantization config's layer lists) need the constituent names + (e.g. `q_proj`) to survive rather than being rewritten to the stacked vLLM name + (`qkv_proj`).""" + return replace(self, orig_to_new_stacked={}) + class AutoWeightsLoader: """ @@ -352,20 +398,19 @@ class AutoWeightsLoader: *, mapper: WeightsMapper | None = None, ) -> set[str]: + # Ignore unexpected biases (typically from GPTQ models) + self.ignore_unexpected_suffixes.append(".bias") + # Many models store quant_config in the base model instead of the causal model. # We look at the causal model's direct children for this reason. modules = (self.module, *self.module.children()) iterator = (m.quant_config for m in modules if hasattr(m, "quant_config")) - quant_config = next(iterator, None) - cache_scale_mapper = ( - quant_config.get_cache_scale_mapper() if quant_config is not None else None - ) - if cache_scale_mapper is not None: - mapper = ( - mapper | cache_scale_mapper - if mapper is not None - else cache_scale_mapper - ) + if quant_config := next(iterator, None): + # Get mappings and ignore prefixes for KV cache quantization scales + mapper = mapper or WeightsMapper() + mapper |= quant_config.get_cache_scale_mapper() + ignore_unexpected_suffixes = quant_config._ignore_unexpected_suffixes + self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes) if mapper is not None: weights = mapper.apply(weights) # filter out weights with first-prefix/substr to skip in name @@ -734,9 +779,7 @@ def maybe_prefix(prefix: str, name: str) -> str: return name if not prefix else f"{prefix}.{name}" -def get_draft_quant_config( - vllm_config: VllmConfig, -) -> QuantizationConfig | None: +def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | None": """Get quantization config for Draft models. Draft models should use their own quantization config instead of the verifier/target diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py index 628186e7598..8efab53db8a 100644 --- a/vllm/model_executor/models/whisper.py +++ b/vllm/model_executor/models/whisper.py @@ -44,7 +44,6 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.whisper_utils import ( ISO639_1_SUPPORTED_LANGS, ) @@ -617,42 +616,6 @@ class WhisperModel(nn.Module): return None return self.encoder(input_features) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".self_attn.qkv_proj", ".self_attn.q_proj", "q"), - (".self_attn.qkv_proj", ".self_attn.k_proj", "k"), - (".self_attn.qkv_proj", ".self_attn.v_proj", "v"), - # MergedColumnParallelLinear uses integer indices (0, 1) - (".encoder_attn.kv_proj", ".encoder_attn.k_proj", 0), - (".encoder_attn.kv_proj", ".encoder_attn.v_proj", 1), - ] - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - class WhisperProcessingInfo(BaseProcessingInfo): def get_hf_config(self) -> WhisperConfig: @@ -808,7 +771,15 @@ class WhisperForConditionalGeneration( } hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."} + orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}, + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), + ".encoder_attn.k_proj": (".encoder_attn.kv_proj", 0), + ".encoder_attn.v_proj": (".encoder_attn.kv_proj", 1), + }, ) # Whisper only supports audio-conditioned generation. From 9e86352c606c61095029f156ae3e4ac2097cf7e5 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Mon, 29 Jun 2026 16:57:26 +0800 Subject: [PATCH 107/138] [CI Failure] Add transformers version check for openai/privacy-filter (#47011) Signed-off-by: wang.yuqi --- tests/models/language/pooling/test_token_classification.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 0f993d965c7..8dc38cf62a0 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -5,6 +5,7 @@ import pytest import torch from transformers import AutoModelForTokenClassification +from tests.models.registry import HF_EXAMPLE_MODELS from tests.models.utils import softmax from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -136,6 +137,9 @@ def test_openai_privacy_filter( model: str, dtype: str, ) -> None: + model_info = HF_EXAMPLE_MODELS.find_hf_info(model) + model_info.check_transformers_version(on_fail="skip") + with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model: vllm_outputs = vllm_model.token_classify(PRIVACY_FILTER_PROMPTS) From 0e207dac784e6b217b8dc1f44ae3985b1f216b50 Mon Sep 17 00:00:00 2001 From: John Date: Mon, 29 Jun 2026 04:59:15 -0400 Subject: [PATCH 108/138] [Bugfix] Transformers backend: apply learned lm_head.bias for tied-embedding models (#46835) Signed-off-by: John Langford Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../layers/vocab_parallel_embedding.py | 17 +++++++---------- .../models/transformers/base.py | 5 +---- .../models/transformers/causal.py | 19 ++++++++++++++++++- 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 61f33591b8c..8d9a7ccbaca 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -542,19 +542,16 @@ class ParallelLMHead(VocabParallelEmbedding): ) self.quant_config = quant_config if bias: - self.bias = Parameter( - torch.empty(self.num_embeddings_per_partition, dtype=params_dtype) - ) - set_weight_attrs( - self.bias, - { - "output_dim": 0, - "weight_loader": self.weight_loader, - }, - ) + self._register_bias() else: self.register_parameter("bias", None) + def _register_bias(self): + data = torch.empty(self.num_embeddings_per_partition, dtype=self.params_dtype) + self.bias = Parameter(data, requires_grad=False) + weight_attrs = dict(output_dim=0, weight_loader=self.weight_loader) + set_weight_attrs(weight=self.bias, weight_attrs=weight_attrs) + def tie_weights(self, embed_tokens: VocabParallelEmbedding): """Tie the weights with word embeddings.""" return self.quant_method.tie_weights(self, embed_tokens) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 4402d180ca0..bcda62918f3 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -640,10 +640,7 @@ class Base( return hidden_states, aux_hidden_states return hidden_states - def load_weights( - self, - weights: Iterable[tuple[str, torch.Tensor]], - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, skip_prefixes=self.skip_prefixes, diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index b6ceb2d6770..01a7e419834 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -16,6 +16,7 @@ # limitations under the License. """Transformers modeling backend mixin for causal language models.""" +from collections.abc import Iterable from typing import TYPE_CHECKING from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -61,6 +62,22 @@ class CausalMixin(VllmModelForTextGeneration): else: self.lm_head = PPMissingLayer() + def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]: + """A thin wrapper around `Base.load_weights` to handle the lm_head bias.""" + + lm_head_bias = set() + + def auto_load_lm_head_bias(weights): + for name, weight in weights: + if name.endswith("lm_head.bias") and self.pp_group.is_last_rank: + self.lm_head._register_bias() + self.lm_head.bias.weight_loader(self.lm_head.bias, weight) + lm_head_bias.add(name) + else: + yield name, weight + + return super().load_weights(auto_load_lm_head_bias(weights)) | lm_head_bias + def compute_logits(self, hidden_states: "torch.Tensor") -> "torch.Tensor | None": - logits = self.logits_processor(self.lm_head, hidden_states) + logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias) return logits From e1861078704b0b091206e83cdd64eaf10b1967ef Mon Sep 17 00:00:00 2001 From: Alden Lobo Date: Mon, 29 Jun 2026 04:12:20 -0500 Subject: [PATCH 109/138] [Bugfix] Use native SiLU activation in CPU fused MoE (#45961) Signed-off-by: Alden Lobo Co-authored-by: Alden Lobo --- vllm/model_executor/layers/fused_moe/cpu_fused_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 868d26e7494..1a0acff058c 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -51,7 +51,7 @@ def _gelu_and_mul( # Uses static methods or standalone functions to avoid instantiating CustomOp # classes, which would call get_current_vllm_config() before config is set. _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = { - MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x), + MoEActivation.SILU: SiluAndMul.forward_native, MoEActivation.SWIGLUOAI: _swigluoai_forward_native, MoEActivation.GELU: _gelu_and_mul, MoEActivation.GELU_TANH: ( From ab132ee98ba14c5d99977b1f83c2d5517c0a1e79 Mon Sep 17 00:00:00 2001 From: soaringk <42689402+soaringk@users.noreply.github.com> Date: Mon, 29 Jun 2026 17:17:54 +0800 Subject: [PATCH 110/138] Fix model info cache for package models (#46567) Signed-off-by: soaringk --- tests/models/test_registry.py | 17 +++++++++++++++++ vllm/model_executor/models/registry.py | 22 ++++++++++++++++++++-- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 0715409abda..7e3ecb372e2 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -20,6 +20,7 @@ from vllm.model_executor.models.registry import ( _SPECULATIVE_DECODING_MODELS, _TEXT_GENERATION_MODELS, ModelRegistry, + _LazyRegisteredModel, ) from vllm.platforms import current_platform @@ -127,6 +128,22 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda): ) +def test_lazy_modelinfo_package_hash_includes_submodules(tmp_path): + package_dir = tmp_path / "model_package" + package_dir.mkdir() + init_file = package_dir / "__init__.py" + init_file.write_text("from .model import Model\n", encoding="utf-8") + model_file = package_dir / "model.py" + model_file.write_text("class Model: pass\n", encoding="utf-8") + + first_hash = _LazyRegisteredModel._get_modelinfo_module_hash(init_file) + + model_file.write_text("class Model:\n supports_pp = True\n", encoding="utf-8") + second_hash = _LazyRegisteredModel._get_modelinfo_module_hash(init_file) + + assert first_hash != second_hash + + def test_hf_registry_coverage(): untested_archs = ( ModelRegistry.get_supported_archs() - HF_EXAMPLE_MODELS.get_supported_archs() diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index efc01033499..ca812c8ee90 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -849,6 +849,25 @@ class _LazyRegisteredModel(_BaseRegisteredModel): cls_name = f"{self.module_name}-{self.class_name}".replace(".", "-") return f"{cls_name}.json" + @staticmethod + def _get_modelinfo_module_hash(model_path: Path) -> str: + if model_path.name == "__init__.py": + # Package entry points often re-export classes implemented in + # submodules, so include the package contents in the cache key. + module_paths = sorted(model_path.parent.rglob("*.py")) + root_path = model_path.parent + else: + module_paths = [model_path] + root_path = model_path.parent + + hasher = safe_hash(b"", usedforsecurity=False) + for path in module_paths: + hasher.update(path.relative_to(root_path).as_posix().encode("utf-8")) + hasher.update(b"\0") + hasher.update(path.read_bytes()) + hasher.update(b"\0") + return hasher.hexdigest() + def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None: try: try: @@ -915,8 +934,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel): module_hash = None if model_path is not None and model_path.exists(): - with open(model_path, "rb") as f: - module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest() + module_hash = self._get_modelinfo_module_hash(model_path) mi = self._load_modelinfo_from_cache(module_hash) if mi is not None: From a4e3cb40d07a1b43f6283cb77d560330b46369a9 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 29 Jun 2026 10:29:09 +0100 Subject: [PATCH 111/138] [mypy] Enable mypy for tests directory (#47018) Signed-off-by: Martin Hickey --- tools/pre_commit/mypy.py | 77 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 76 insertions(+), 1 deletion(-) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index ccbce700441..32d7f45f318 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -25,6 +25,81 @@ import regex as re # from "skip" to "silent", remove its directory from SEPARATE_GROUPS. SEPARATE_GROUPS = [ "tests", + "tests/benchmarks", + "tests/compile/correctness_e2e", + "tests/config", + "tests/compile", + "tests/compile/fullgraph", + "tests/compile/fusions_e2e", + "tests/compile/passes", + "tests/distributed", + "tests/entrypoints/anthropic", + "tests/entrypoints/generate", + "tests/entrypoints/llm", + "tests/entrypoints/multimodal", + "tests/entrypoints/openai", + "tests/entrypoints/pooling", + "tests/entrypoints/serve", + "tests/entrypoints/speech_to_text", + "tests/entrypoints/tool_parsers", + "tests/entrypoints/unit_tests", + "tests/entrypoints/weight_transfer", + "tests/kernels", + "tests/kernels/attention", + "tests/kernels/core", + "tests/kernels/helion", + "tests/kernels/mamba", + "tests/kernels/moe", + "tests/kernels/quantization", + "tests/lora", + "tests/model_executor", + "tests/model_executor/layers", + "tests/model_executor/model_loader", + "tests/models", + "tests/models/test_initialization.py", + "tests/models/language", + "tests/models/multimodal", + "tests/models/quantization", + "tests/multimodal", + "tests/parser", + "tests/plugins_tests/gguf", + "tests/plugins_tests/lora_resolvers", + "tests/plugins/bge_m3_sparse_plugin", + "tests/plugins/prithvi_io_processor_plugin", + "tests/plugins/vllm_add_dummy_platform", + "tests/plugins/vllm_add_dummy_stat_logger", + "tests/plugins_tests", + "tests/quantization", + "tests/reasoning", + "tests/renderers", + "tests/samplers", + "tests/spec_decode", + "tests/tokenizers_", + "tests/tool_parsers", + "tests/tool_use", + "tests/transformers_utils", + "tests/utils_", + "tests/v1", + "tests/v1/attention", + "tests/v1/core", + "tests/v1/cudagraph", + "tests/v1/determinism", + "tests/v1/distributed", + "tests/v1/e2e", + "tests/v1/ec_connector", + "tests/v1/engine", + "tests/v1/executor", + "tests/v1/kv_connector", + "tests/v1/kv_offload", + "tests/v1/logits_processors", + "tests/v1/metrics", + "tests/v1/sample", + "tests/v1/shutdown", + "tests/v1/simple_kv_offload", + "tests/v1/spec_decode", + "tests/v1/streaming_input", + "tests/v1/structured_output", + "tests/v1/worker", ] # TODO(woosuk): Include the code from Megatron and HuggingFace. @@ -57,7 +132,7 @@ def group_files(changed_files: list[str]) -> dict[str, list[str]]: file_groups[directory].append(changed_file) break else: - if changed_file.startswith("vllm/"): + if changed_file.startswith(("vllm/", "tests/")): file_groups[""].append(changed_file) return file_groups From eddfd4cf219359296758272ca736d38cb2c327b1 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Mon, 29 Jun 2026 12:10:07 +0200 Subject: [PATCH 112/138] [Perf][2/N] Expand Triton kernel warmup coverage, Qwen (#46750) Signed-off-by: LopezCastroRoberto --- vllm/model_executor/warmup/kernel_warmup.py | 3 + .../warmup/qwen_triton_warmup.py | 386 ++++++++++++++++++ 2 files changed, 389 insertions(+) create mode 100644 vllm/model_executor/warmup/qwen_triton_warmup.py diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 754270e6525..7edbff4d4a6 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -24,6 +24,7 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( deepseek_v4_sparse_mla_attention_warmup, flashinfer_sparse_mla_decode_autotune_warmup, ) +from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -40,6 +41,8 @@ def kernel_warmup(worker: "Worker"): minimax_m3_msa_warmup, ) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) + # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder # layer per token; warm them across token sizes first so the first real # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside). diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py new file mode 100644 index 00000000000..62e94f93c09 --- /dev/null +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up Qwen Triton kernels from the loaded model's compile keys.""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import ( + fused_post_conv_prep, +) +from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import ( + fused_sigmoid_gating_delta_rule_update, +) +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn, +) +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID +from vllm.v1.worker.block_table import BlockTable +from vllm.v1.worker.utils import _zero_kv_blocks_kernel + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +logger = init_logger(__name__) + +_QWEN_MODEL_TYPES = frozenset( + { + "qwen3_next", + "qwen3_5", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + } +) + +_ZERO_KV_N_BLOCKS = (1, 2) + +_SLOT_MAPPING_KV_BLOCK_SIZE = 16 +_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1 +_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3) + +# Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L. +_FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16) + + +@dataclass(frozen=True) +class _ZeroKvWarmupConfig: + page_size_el: int + block_size: int + n_segs: int + + +@dataclass(frozen=True) +class _QwenGDNWarmupConfig: + h: int + hv: int + k: int + v: int + conv_kernel_size: int + conv_state: torch.Tensor + conv_dtype: torch.dtype + a_log: torch.Tensor + dt_bias: torch.Tensor + state_stride_token: int + state_dtype: torch.dtype + + @property + def conv_dim(self) -> int: + return 2 * self.h * self.k + self.hv * self.v + + +def _is_non_empty_tensor(value: object) -> bool: + return isinstance(value, torch.Tensor) and value.numel() > 0 + + +def _is_qwen_gdn_layer(module: object) -> bool: + return all( + hasattr(module, attr) + for attr in ( + "num_k_heads", + "num_v_heads", + "head_k_dim", + "head_v_dim", + "conv_kernel_size", + "tp_size", + "kv_cache", + "A_log", + "dt_bias", + ) + ) + + +def _iter_qwen_gdn_layers(static_forward_context: object): + if not isinstance(static_forward_context, dict): + return + + for module in static_forward_context.values(): + if _is_qwen_gdn_layer(module): + yield module + + +def _split_qwen_gdn_cache(kv_cache: object) -> tuple[torch.Tensor, torch.Tensor] | None: + if isinstance(kv_cache, (list, tuple)) and len(kv_cache) >= 2: + conv_cache, ssm_state = kv_cache[:2] + if _is_non_empty_tensor(conv_cache) and _is_non_empty_tensor(ssm_state): + return conv_cache, ssm_state + + if isinstance(kv_cache, torch.Tensor) and kv_cache.size(0) >= 2: + conv_cache = kv_cache[0] + ssm_state = kv_cache[1] + if _is_non_empty_tensor(conv_cache) and _is_non_empty_tensor(ssm_state): + return conv_cache, ssm_state + return None + + +def _qwen_gdn_warmup_config( + static_forward_context: object, +) -> _QwenGDNWarmupConfig | None: + found_layer = False + for layer in _iter_qwen_gdn_layers(static_forward_context): + found_layer = True + cache_tensors = _split_qwen_gdn_cache(getattr(layer, "kv_cache", None)) + if cache_tensors is None: + continue + + conv_cache, ssm_state = cache_tensors + conv_state = ( + conv_cache if is_conv_state_dim_first() else conv_cache.transpose(-1, -2) + ) + tp_size = int(layer.tp_size) + h = int(layer.num_k_heads) // tp_size + hv = int(layer.num_v_heads) // tp_size + + return _QwenGDNWarmupConfig( + h=h, + hv=hv, + k=int(layer.head_k_dim), + v=int(layer.head_v_dim), + conv_kernel_size=int(layer.conv_kernel_size), + conv_state=conv_state, + conv_dtype=conv_state.dtype, + a_log=layer.A_log, + dt_bias=layer.dt_bias, + state_stride_token=int(ssm_state.stride(0)), + state_dtype=ssm_state.dtype, + ) + + if found_layer: + logger.info("Skipping Qwen GDN Triton warmup: no bound Qwen GDN cache found.") + else: + logger.info("Skipping Qwen GDN Triton warmup: no Qwen GDN layer found.") + return None + + +def _get_kv_block_zeroer(runner: object) -> object | None: + zeroer = getattr(runner, "kv_block_zeroer", None) + if zeroer is None: + zeroer = getattr(runner, "_kv_block_zeroer", None) + return zeroer + + +def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None: + zeroer = _get_kv_block_zeroer(runner) + meta = getattr(zeroer, "_meta", None) + if meta is None: + return None + + _, page_size_el, block_size, n_segs = meta + return _ZeroKvWarmupConfig( + page_size_el=int(page_size_el), + block_size=int(block_size), + n_segs=int(n_segs), + ) + + +def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool: + zeroer = _get_kv_block_zeroer(runner) + zero_block_ids = getattr(zeroer, "zero_block_ids", None) + if not callable(zero_block_ids): + return False + + for n_blocks in _ZERO_KV_N_BLOCKS: + zero_block_ids(list(range(n_blocks))) + return True + + +def _warm_zero_kv_blocks_kernel( + device: torch.device, config: _ZeroKvWarmupConfig +) -> None: + max_n_blocks = max(_ZERO_KV_N_BLOCKS) + scratch = torch.empty( + max_n_blocks * config.page_size_el, + dtype=torch.int32, + device=device, + ) + seg_addrs = torch.tensor( + [scratch.data_ptr()] * config.n_segs, + dtype=torch.uint64, + device=device, + ) + + for n_blocks in _ZERO_KV_N_BLOCKS: + block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device) + grid = (n_blocks * config.n_segs * (config.page_size_el // config.block_size),) + _zero_kv_blocks_kernel[grid]( + seg_addrs, + block_ids, + n_blocks, + N_SEGS=config.n_segs, + PAGE_SIZE_EL=config.page_size_el, + BLOCK_SIZE=config.block_size, + num_warps=4, + num_stages=3, + ) + + +def _warm_compute_slot_mapping_kernel(device: torch.device) -> None: + # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny. + num_tokens = 1 + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + + for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES: + # Use BlockTable so the JIT key matches the production slot-mapping call. + block_table = BlockTable( + block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, + max_num_reqs=1, + max_num_blocks_per_req=block_table_stride, + max_num_batched_tokens=num_tokens, + pin_memory=False, + device=device, + kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, + cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE, + ) + block_table.add_row(list(range(block_table_stride)), 0) + block_table.commit_block_table(num_reqs=1) + block_table.compute_slot_mapping(1, query_start_loc, positions) + + +def _warm_causal_conv1d_fwd_kernel( + device: torch.device, config: _QwenGDNWarmupConfig +) -> None: + x_storage = torch.empty( + (1, config.conv_dim), dtype=config.conv_dtype, device=device + ) + x = x_storage.t() + weight = torch.empty( + (config.conv_dim, config.conv_kernel_size), + dtype=config.conv_dtype, + device=device, + ) + cache_indices = torch.full((1,), NULL_BLOCK_ID, dtype=torch.int32, device=device) + has_initial_state = torch.empty(1, dtype=torch.bool, device=device) + query_start_loc = torch.tensor([0, 1], dtype=torch.int32, device=device) + + causal_conv1d_fn( + x, + weight, + None, + config.conv_state, + query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation="silu", + pad_slot_id=PAD_SLOT_ID, + null_block_id=NULL_BLOCK_ID, + metadata=None, + validate_data=False, + ) + + +def _warm_fused_post_conv_kernel( + device: torch.device, config: _QwenGDNWarmupConfig +) -> None: + qkv_dim = 2 * config.h * config.k + config.hv * config.v + for length in _FLA_POST_CONV_WARMUP_LENGTHS: + conv_output = torch.empty( + (length, qkv_dim), dtype=config.conv_dtype, device=device + ) + a = torch.empty((length, config.hv), dtype=config.conv_dtype, device=device) + b = torch.empty_like(a) + + fused_post_conv_prep( + conv_output, + a, + b, + config.a_log, + config.dt_bias, + config.h, + config.k, + config.v, + apply_l2norm=True, + output_g_exp=False, + ) + + +def _warm_fused_sigmoid_gating_delta_rule_update_kernel( + device: torch.device, + config: _QwenGDNWarmupConfig, +) -> None: + q = torch.empty((1, 1, config.h, config.k), dtype=config.conv_dtype, device=device) + k = torch.empty_like(q) + v = torch.empty((1, 1, config.hv, config.v), dtype=config.conv_dtype, device=device) + a = torch.empty((1, 1, config.hv), dtype=config.conv_dtype, device=device) + b = torch.empty_like(a) + state = torch.empty( + (1, config.state_stride_token), + dtype=config.state_dtype, + device=device, + ) + cu_seqlens = torch.tensor([0, 1], dtype=torch.int32, device=device) + ssm_state_indices = torch.empty((1, 1), dtype=torch.int32, device=device) + ssm_state_indices.zero_() + + fused_sigmoid_gating_delta_rule_update( + A_log=config.a_log, + a=a, + b=b, + dt_bias=config.dt_bias, + q=q, + k=k, + v=v, + beta=1.0, + threshold=20.0, + initial_state=state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + use_qk_l2norm_in_kernel=True, + is_kda=False, + ) + + +def _synchronize_device(device: torch.device) -> None: + if device.type == "cuda": + torch.accelerator.synchronize(device) + + +@torch.inference_mode() +def qwen_triton_warmup( + runner: "GPUModelRunner", + model_config: object, +) -> None: + """Warm Qwen Triton kernels reported by the JIT monitor.""" + if runner.is_pooling_model: + return + + hf_text_config = getattr(model_config, "hf_text_config", None) + hf_config = getattr(model_config, "hf_config", None) + model_type = None + for config in (hf_text_config, hf_config): + model_type = getattr(config, "model_type", None) + if model_type is not None: + model_type = str(model_type) + break + if model_type not in _QWEN_MODEL_TYPES: + return + + device = getattr(runner, "device", torch.device("cuda")) + logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type) + + zero_config = _zero_kv_warmup_config(runner) + if _warm_zero_kv_blocks_with_runner_zeroer(runner): + pass + elif zero_config is not None: + _warm_zero_kv_blocks_kernel(device, zero_config) + else: + logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.") + + _warm_compute_slot_mapping_kernel(device) + _synchronize_device(device) + + compilation_config = getattr(runner, "compilation_config", None) + static_forward_context = getattr(compilation_config, "static_forward_context", None) + gdn_config = _qwen_gdn_warmup_config(static_forward_context) + if gdn_config is None: + return + + _warm_causal_conv1d_fwd_kernel(device, gdn_config) + _warm_fused_post_conv_kernel(device, gdn_config) + _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config) + _synchronize_device(device) From 3483240b7ea3d4372b6c79369ea36617f8b1fbb2 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Mon, 29 Jun 2026 18:18:53 +0800 Subject: [PATCH 113/138] [Frontend] Consolidate scale out entrypoints (#44512) Signed-off-by: wang.yuqi --- .buildkite/test-amd.yaml | 4 + .buildkite/test_areas/entrypoints.yaml | 2 + .buildkite/test_areas/rust_frontend.yaml | 4 +- docs/examples/README.md | 3 +- docs/serving/online_serving/README.md | 4 +- .../serve/disagg => examples}/__init__.py | 0 .../render => examples/scale_out}/__init__.py | 0 .../example_mm_serve.py | 0 .../token_generation_client.py | 0 .../openai/chat_completion/test_chat_error.py | 12 +- .../completion/test_completion_error.py | 12 +- .../openai/test_render_token_offsets.py | 4 +- .../entrypoints/scale_out}/__init__.py | 0 .../scale_out/derender/__init__.py | 0 .../derender}/test_derender.py | 0 .../entrypoints/scale_out/render/__init__.py | 0 .../render/test_launch_render.py | 0 .../render/test_render.py | 0 .../render/test_render_multimodal.py | 0 .../scale_out/token_in_token_out/__init__.py | 0 .../test_generate_stream.py | 4 +- .../token_in_token_out}/test_mm_serde.py | 9 +- .../token_in_token_out}/test_protocol.py | 4 +- .../test_return_routed_experts.py | 0 .../test_serving_multimodal_tokens.py | 0 .../test_serving_tokens.py | 0 .../test_tokens_logprobs.py | 2 +- tests/renderers/test_token_offsets.py | 2 +- vllm/entrypoints/generate/api_router.py | 15 -- vllm/entrypoints/openai/api_server.py | 33 +-- vllm/entrypoints/scale_out/__init__.py | 0 .../scale_out/derender/__init__.py | 0 .../scale_out/derender/api_router.py | 74 +++++++ .../entrypoints/scale_out/derender/serving.py | 202 ++++++++++++++++++ vllm/entrypoints/scale_out/factories.py | 78 +++++++ .../{serve => scale_out}/render/__init__.py | 0 .../{serve => scale_out}/render/api_router.py | 73 +------ .../{serve => scale_out}/render/serving.py | 136 +----------- .../scale_out/token_in_token_out/__init__.py | 0 .../token_in_token_out}/api_router.py | 13 +- .../token_in_token_out}/mm_serde.py | 0 .../token_in_token_out}/protocol.py | 0 .../token_in_token_out}/serving.py | 17 +- vllm/entrypoints/serve/engine/typing.py | 2 +- vllm/renderers/online_derenderer.py | 2 +- 45 files changed, 422 insertions(+), 289 deletions(-) rename {tests/entrypoints/serve/disagg => examples}/__init__.py (100%) rename {tests/entrypoints/serve/render => examples/scale_out}/__init__.py (100%) rename examples/{disaggregated/disaggregated_serving => scale_out}/example_mm_serve.py (100%) rename examples/{generate => scale_out}/token_generation_client.py (100%) rename {vllm/entrypoints/serve/disagg => tests/entrypoints/scale_out}/__init__.py (100%) create mode 100644 tests/entrypoints/scale_out/derender/__init__.py rename tests/entrypoints/{serve/render => scale_out/derender}/test_derender.py (100%) create mode 100644 tests/entrypoints/scale_out/render/__init__.py rename tests/entrypoints/{serve => scale_out}/render/test_launch_render.py (100%) rename tests/entrypoints/{serve => scale_out}/render/test_render.py (100%) rename tests/entrypoints/{serve => scale_out}/render/test_render_multimodal.py (100%) create mode 100644 tests/entrypoints/scale_out/token_in_token_out/__init__.py rename tests/entrypoints/{serve/disagg => scale_out/token_in_token_out}/test_generate_stream.py (99%) rename tests/entrypoints/{openai => scale_out/token_in_token_out}/test_mm_serde.py (94%) rename tests/entrypoints/{serve/disagg => scale_out/token_in_token_out}/test_protocol.py (95%) rename tests/entrypoints/{serve/disagg => scale_out/token_in_token_out}/test_return_routed_experts.py (100%) rename tests/entrypoints/{serve/disagg => scale_out/token_in_token_out}/test_serving_multimodal_tokens.py (100%) rename tests/entrypoints/{serve/disagg => scale_out/token_in_token_out}/test_serving_tokens.py (100%) rename tests/entrypoints/{serve/disagg => scale_out/token_in_token_out}/test_tokens_logprobs.py (92%) create mode 100644 vllm/entrypoints/scale_out/__init__.py create mode 100644 vllm/entrypoints/scale_out/derender/__init__.py create mode 100644 vllm/entrypoints/scale_out/derender/api_router.py create mode 100644 vllm/entrypoints/scale_out/derender/serving.py create mode 100644 vllm/entrypoints/scale_out/factories.py rename vllm/entrypoints/{serve => scale_out}/render/__init__.py (100%) rename vllm/entrypoints/{serve => scale_out}/render/api_router.py (50%) rename vllm/entrypoints/{serve => scale_out}/render/serving.py (66%) create mode 100644 vllm/entrypoints/scale_out/token_in_token_out/__init__.py rename vllm/entrypoints/{serve/disagg => scale_out/token_in_token_out}/api_router.py (96%) rename vllm/entrypoints/{serve/disagg => scale_out/token_in_token_out}/mm_serde.py (100%) rename vllm/entrypoints/{serve/disagg => scale_out/token_in_token_out}/protocol.py (100%) rename vllm/entrypoints/{serve/disagg => scale_out/token_in_token_out}/serving.py (99%) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 7521901c9a6..ea76ae1c37f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -845,10 +845,12 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/serve + - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc + - pytest -v -s entrypoints/scale_out - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD timeout_in_minutes: 180 @@ -2559,10 +2561,12 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/serve + - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc + - pytest -v -s entrypoints/scale_out - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD timeout_in_minutes: 180 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index f6307f097d9..d95b7e0d008 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -40,10 +40,12 @@ steps: source_file_dependencies: - vllm/ - tests/entrypoints/serve + - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc + - pytest -v -s entrypoints/scale_out mirror: amd: device: mi325_1 diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 5ea0f7ef77c..adb27c4a049 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -46,7 +46,7 @@ steps: - vllm/v1/engine/ - tests/utils.py # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py - - tests/entrypoints/serve/disagg/test_serving_tokens.py + - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py - tests/entrypoints/serve/instrumentator/test_basic.py - tests/entrypoints/serve/instrumentator/test_metrics.py # - tests/entrypoints/serve/dev/test_sleep.py @@ -55,7 +55,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" - - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" + - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" # - pytest -v -s entrypoints/serve/dev/test_sleep.py diff --git a/docs/examples/README.md b/docs/examples/README.md index 5569db9119c..a9a127a4d5d 100644 --- a/docs/examples/README.md +++ b/docs/examples/README.md @@ -13,5 +13,6 @@ vLLM's examples are organized into the following categories: - **[`rl/`](../../examples/rl)** โ€“ Reinforcement learning examples. - **[`deployment/`](../../examples/deployment)** โ€“ Examples for deploying vLLM in production. - **[`ray_serving/`](../../examples/ray_serving)** โ€“ Scalable serving using Ray. -- **[`disaggregated/`](../../examples/disaggregated)** โ€“ Examples for disaggregated serving (separate prefill and decode), including various kv cache connectors (LMCache, Mooncake, FlexKV, P2P NCCL) and failure recovery. +- **[`disaggregated/`](../../examples/disaggregated)** โ€“ Examples for Disaggregated P/D (Prefill/Decoding) inference, including various kv cache connectors (LMCache, Mooncake, FlexKV, P2P NCCL) and failure recovery. +- **[`scale_out/`](../../examples/scale_out)** โ€“ Examples for Token In <> Token Out API Server. - **[`observability/`](../../examples/observability)** โ€“ Metrics, logging, tracing (OpenTelemetry), and dashboards (Grafana, Perses). diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 40fc8b7c426..60476fa5edb 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -119,9 +119,9 @@ For further details on profiling vLLM, please refer to [this page](../../contrib - `/ping` - SageMaker health check - `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints) -## Disaggregated Everything +## Scale-Out APIs -### Tokens IN <> Tokens OUT +### Tokens IN <> Tokens OUT APIs - `/inference/v1/generate` - Generate completions - `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set) diff --git a/tests/entrypoints/serve/disagg/__init__.py b/examples/__init__.py similarity index 100% rename from tests/entrypoints/serve/disagg/__init__.py rename to examples/__init__.py diff --git a/tests/entrypoints/serve/render/__init__.py b/examples/scale_out/__init__.py similarity index 100% rename from tests/entrypoints/serve/render/__init__.py rename to examples/scale_out/__init__.py diff --git a/examples/disaggregated/disaggregated_serving/example_mm_serve.py b/examples/scale_out/example_mm_serve.py similarity index 100% rename from examples/disaggregated/disaggregated_serving/example_mm_serve.py rename to examples/scale_out/example_mm_serve.py diff --git a/examples/generate/token_generation_client.py b/examples/scale_out/token_generation_client.py similarity index 100% rename from examples/generate/token_generation_client.py rename to examples/scale_out/token_generation_client.py diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index 3eea57d3f53..4b6be87ae5c 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -17,10 +17,9 @@ from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.render.serving import ServingRender +from vllm.entrypoints.scale_out.render.serving import ServingRender from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.hf import HfRenderer -from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM @@ -206,15 +205,8 @@ def _build_serving_render(engine: AsyncLLM) -> ServingRender: chat_template=None, chat_template_content_format="auto", ) - online_derenderer = OnlineDerenderer( - model_config=engine.model_config, - renderer=engine.renderer, - request_logger=None, - chat_template=None, - chat_template_content_format="auto", - ) - serving_render = ServingRender(models, online_renderer, online_derenderer) + serving_render = ServingRender(models, online_renderer) async def _fake_preprocess_chat(*args, **kwargs): # return conversation, engine_inputs diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 9d2fedae361..062c3e7583a 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -14,10 +14,9 @@ from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.render.serving import ServingRender +from vllm.entrypoints.scale_out.render.serving import ServingRender from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.hf import HfRenderer -from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM @@ -191,15 +190,8 @@ def _build_serving_render(engine: AsyncLLM) -> ServingRender: chat_template=None, chat_template_content_format="auto", ) - online_derenderer = OnlineDerenderer( - model_config=engine.model_config, - renderer=engine.renderer, - request_logger=None, - chat_template=None, - chat_template_content_format="auto", - ) - serving_render = ServingRender(models, online_renderer, online_derenderer) + serving_render = ServingRender(models, online_renderer) async def _fake_preprocess_chat(*args, **kwargs): # return conversation, engine_inputs diff --git a/tests/entrypoints/openai/test_render_token_offsets.py b/tests/entrypoints/openai/test_render_token_offsets.py index f7653ab66fc..a2e66b7bd6c 100644 --- a/tests/entrypoints/openai/test_render_token_offsets.py +++ b/tests/entrypoints/openai/test_render_token_offsets.py @@ -3,7 +3,7 @@ """Unit tests for the token-offsets request/response protocol wiring: the request flag flowing into ``TokenizeParams`` and the ``GenerateRequest`` serialization boundary. End-to-end behavior is covered by -``tests/entrypoints/serve/render/test_render.py``; plain Pydantic field +``tests/entrypoints/scale_out/render/test_render.py``; plain Pydantic field storage is not retested here. """ @@ -12,7 +12,7 @@ from unittest.mock import Mock from vllm.config import ModelConfig from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest -from vllm.entrypoints.serve.disagg.protocol import GenerateRequest +from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateRequest from vllm.sampling_params import SamplingParams diff --git a/vllm/entrypoints/serve/disagg/__init__.py b/tests/entrypoints/scale_out/__init__.py similarity index 100% rename from vllm/entrypoints/serve/disagg/__init__.py rename to tests/entrypoints/scale_out/__init__.py diff --git a/tests/entrypoints/scale_out/derender/__init__.py b/tests/entrypoints/scale_out/derender/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/entrypoints/serve/render/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py similarity index 100% rename from tests/entrypoints/serve/render/test_derender.py rename to tests/entrypoints/scale_out/derender/test_derender.py diff --git a/tests/entrypoints/scale_out/render/__init__.py b/tests/entrypoints/scale_out/render/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/entrypoints/serve/render/test_launch_render.py b/tests/entrypoints/scale_out/render/test_launch_render.py similarity index 100% rename from tests/entrypoints/serve/render/test_launch_render.py rename to tests/entrypoints/scale_out/render/test_launch_render.py diff --git a/tests/entrypoints/serve/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py similarity index 100% rename from tests/entrypoints/serve/render/test_render.py rename to tests/entrypoints/scale_out/render/test_render.py diff --git a/tests/entrypoints/serve/render/test_render_multimodal.py b/tests/entrypoints/scale_out/render/test_render_multimodal.py similarity index 100% rename from tests/entrypoints/serve/render/test_render_multimodal.py rename to tests/entrypoints/scale_out/render/test_render_multimodal.py diff --git a/tests/entrypoints/scale_out/token_in_token_out/__init__.py b/tests/entrypoints/scale_out/token_in_token_out/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py similarity index 99% rename from tests/entrypoints/serve/disagg/test_generate_stream.py rename to tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index a31655e4307..ce3100f196c 100644 --- a/tests/entrypoints/serve/disagg/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -12,11 +12,11 @@ from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.engine.protocol import StreamOptions from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.disagg.protocol import ( +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, GenerateResponse, ) -from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.logprobs import Logprob from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers import renderer_from_config diff --git a/tests/entrypoints/openai/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py similarity index 94% rename from tests/entrypoints/openai/test_mm_serde.py rename to tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index c568d822e1c..d24436bbd4b 100644 --- a/tests/entrypoints/openai/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -1,14 +1,17 @@ # 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.""" +""" +Roundtrip tests for multimodal serde used by the +token_in_token_out generate endpoint. +""" import torch -from vllm.entrypoints.serve.disagg.mm_serde import ( +from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import ( decode_mm_kwargs_item, encode_mm_kwargs_item, ) -from vllm.entrypoints.serve.disagg.protocol import ( +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( MultiModalFeatures, PlaceholderRangeInfo, ) diff --git a/tests/entrypoints/serve/disagg/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py similarity index 95% rename from tests/entrypoints/serve/disagg/test_protocol.py rename to tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index 414fc2a2612..674ce18b7f3 100644 --- a/tests/entrypoints/serve/disagg/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the disagg request/response protocol. +"""Unit tests for the token_in_token_out request/response protocol. These tests intentionally avoid spinning up a server โ€” they exercise the pydantic validators on ``GenerateRequest`` directly so they run fast and @@ -9,7 +9,7 @@ fail loudly if the validator semantics ever drift. import json -from vllm.entrypoints.serve.disagg.protocol import GenerateRequest +from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateRequest from vllm.sampling_params import SamplingParams diff --git a/tests/entrypoints/serve/disagg/test_return_routed_experts.py b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py similarity index 100% rename from tests/entrypoints/serve/disagg/test_return_routed_experts.py rename to tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py diff --git a/tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py similarity index 100% rename from tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py rename to tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py diff --git a/tests/entrypoints/serve/disagg/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py similarity index 100% rename from tests/entrypoints/serve/disagg/test_serving_tokens.py rename to tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py diff --git a/tests/entrypoints/serve/disagg/test_tokens_logprobs.py b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py similarity index 92% rename from tests/entrypoints/serve/disagg/test_tokens_logprobs.py rename to tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py index 844dd24d541..80f08078da2 100644 --- a/tests/entrypoints/serve/disagg/test_tokens_logprobs.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.logprobs import Logprob diff --git a/tests/renderers/test_token_offsets.py b/tests/renderers/test_token_offsets.py index ab881782659..f973e8610bc 100644 --- a/tests/renderers/test_token_offsets.py +++ b/tests/renderers/test_token_offsets.py @@ -5,7 +5,7 @@ These exercise ``_tokenize_prompt`` (offset extraction + capability/MM gating) and the ``_tokenize_prompt -> _process_tokens -> TokensInput`` forwarding chain. Endpoint-level coverage lives in -``tests/entrypoints/serve/render/test_render.py``. +``tests/entrypoints/scale_out/render/test_render.py``. """ import pytest diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 38ecdec5ce2..5a26e475b06 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -66,7 +66,6 @@ async def init_generate_state( from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses - from vllm.entrypoints.serve.disagg.serving import ServingTokens from vllm.entrypoints.serve.utils.fingerprint import set_default_fingerprint_mode # Applied before any serving class is constructed so that each one picks @@ -175,20 +174,6 @@ async def init_generate_state( if "generate" in supported_tasks else None ) - state.serving_tokens = ( - ServingTokens( - engine_client, - state.openai_serving_models, - state.online_renderer, - request_logger=request_logger, - return_tokens_as_token_ids=args.return_tokens_as_token_ids, - enable_prompt_tokens_details=args.enable_prompt_tokens_details, - enable_log_outputs=args.enable_log_outputs, - force_no_detokenize=args.tokens_only, - ) - if "generate" in supported_tasks - else None - ) from .generative_scoring.serving import ServingGenerativeScoring diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 9fc4560adbe..6ae6dd70abd 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -32,7 +32,6 @@ from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware -from vllm.entrypoints.serve.render.serving import ServingRender from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( @@ -208,12 +207,6 @@ def build_app( register_generate_api_routers(app) - from vllm.entrypoints.serve.disagg.api_router import ( - attach_router as attach_disagg_router, - ) - - attach_disagg_router(app) - from vllm.entrypoints.serve.elastic_ep.api_router import ( attach_router as elastic_ep_attach_router, ) @@ -221,11 +214,9 @@ def build_app( elastic_ep_attach_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, - ) + from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers - attach_render_router(app) + register_scale_out_api_routers(app, supported_tasks) if "transcription" in supported_tasks or "realtime" in supported_tasks: from vllm.entrypoints.speech_to_text.factories import ( @@ -401,12 +392,6 @@ async def init_app_state( default_chat_template_kwargs=args.default_chat_template_kwargs, trust_request_chat_template=args.trust_request_chat_template, ) - state.serving_render = ServingRender( - state.openai_serving_models, - state.online_renderer, - state.online_derenderer, - request_logger=request_logger, - ) if "generate" in supported_tasks: from vllm.entrypoints.generate.api_router import init_generate_state @@ -415,6 +400,10 @@ async def init_app_state( engine_client, state, args, request_logger, supported_tasks ) + from vllm.entrypoints.scale_out.factories import init_scale_out_state + + init_scale_out_state(state, args, engine_client, request_logger) + if "transcription" in supported_tasks or "realtime" in supported_tasks: from vllm.entrypoints.speech_to_text.factories import init_speech_to_text_state @@ -505,12 +494,10 @@ async def init_render_app_state( default_chat_template_kwargs=args.default_chat_template_kwargs, trust_request_chat_template=args.trust_request_chat_template, ) - state.serving_render = ServingRender( - model_registry, - state.online_renderer, - state.online_derenderer, - request_logger=request_logger, - ) + + from vllm.entrypoints.scale_out.factories import init_render_state + + init_render_state(state, request_logger) state.vllm_config = vllm_config # Disable stats logging โ€” there is no engine to poll. diff --git a/vllm/entrypoints/scale_out/__init__.py b/vllm/entrypoints/scale_out/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/scale_out/derender/__init__.py b/vllm/entrypoints/scale_out/derender/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/scale_out/derender/api_router.py b/vllm/entrypoints/scale_out/derender/api_router.py new file mode 100644 index 00000000000..3f88d51f0a9 --- /dev/null +++ b/vllm/entrypoints/scale_out/derender/api_router.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from http import HTTPStatus + +from fastapi import APIRouter, Depends, Request +from fastapi.responses import JSONResponse + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse +from vllm.entrypoints.openai.completion.protocol import CompletionResponse +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.utils.api_utils import validate_json_request +from vllm.logger import init_logger + +from ..token_in_token_out.protocol import ( + DerenderChatRequest, + DerenderCompletionRequest, +) +from .serving import ServingDerender + +logger = init_logger(__name__) + +router = APIRouter() + + +def derender(request: Request) -> ServingDerender | None: + return getattr(request.app.state, "serving_derender", None) + + +@router.post( + "/v1/chat/completions/derender", + dependencies=[Depends(validate_json_request)], + response_model=ChatCompletionResponse, + responses={ + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +async def derender_chat_completion(request: DerenderChatRequest, raw_request: Request): + handler = derender(raw_request) + if handler is None: + raise NotImplementedError( + "The model does not support Chat Completions Derender API" + ) + + result = await handler.derender_chat_response(request) + + if isinstance(result, ErrorResponse): + return JSONResponse(content=result.model_dump(), status_code=result.error.code) + + return JSONResponse(content=result.model_dump()) + + +@router.post( + "/v1/completions/derender", + dependencies=[Depends(validate_json_request)], + response_model=CompletionResponse, + responses={ + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +async def derender_completion(request: DerenderCompletionRequest, raw_request: Request): + handler = derender(raw_request) + if handler is None: + raise NotImplementedError("The model does not support Completions Derender API") + + result = await handler.derender_completion_response(request) + + if isinstance(result, ErrorResponse): + return JSONResponse(content=result.model_dump(), status_code=result.error.code) + + return JSONResponse(content=result.model_dump()) diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py new file mode 100644 index 00000000000..e125007a549 --- /dev/null +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time +from typing import cast + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse +from vllm.entrypoints.openai.completion.protocol import CompletionResponse +from vllm.entrypoints.openai.engine.protocol import ( + ErrorResponse, + UsageInfo, +) +from vllm.entrypoints.openai.models.serving import ( + OpenAIModelRegistry, + OpenAIServingModels, +) +from vllm.entrypoints.serve.engine.serving import BaseServing +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.inputs import ( + EngineInput, + MultiModalHashes, + MultiModalInput, + MultiModalPlaceholders, +) +from vllm.logger import init_logger +from vllm.renderers.online_derenderer import OnlineDerenderer + +from ..token_in_token_out.mm_serde import encode_mm_kwargs_item +from ..token_in_token_out.protocol import ( + DerenderChatRequest, + DerenderCompletionRequest, + MultiModalFeatures, + PlaceholderRangeInfo, +) + +logger = init_logger(__name__) + + +class ServingDerender(BaseServing): + def __init__( + self, + models: OpenAIServingModels | OpenAIModelRegistry, + online_derenderer: "OnlineDerenderer", + *, + request_logger: RequestLogger | None = None, + ) -> None: + super().__init__( + models=models, + model_config=models.model_config, + request_logger=request_logger, + ) + + self.online_derenderer = online_derenderer + + async def derender_chat_response( + self, + request: DerenderChatRequest, + ) -> ChatCompletionResponse | ErrorResponse: + """Postprocess a GenerateResponse into a ChatCompletionResponse. + + Non-streaming only: expects the complete GenerateResponse with all + token IDs present. Uses ``parser.parse()`` for one-shot extraction. + + When ``request.chat_request`` is provided, the parser splits the + output into (reasoning, content, tool_calls). Otherwise falls + back to plain detokenization. + """ + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + return error_check_ret + + try: + choices = await self.online_derenderer.derender_chat( + request.generate_response, request.chat_request + ) + except ValueError as exc: + return self.create_error_response(str(exc)) + + prompt_tokens = ( + request.prompt_tokens if request.prompt_tokens is not None else 0 + ) + gen = request.generate_response + completion_tokens = sum(len(ch.token_ids) for ch in gen.choices if ch.token_ids) + usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + logger.debug( + "derender_chat request_id=%s model=%s choices=%d completion_tokens=%d", + gen.request_id, + request.model, + len(choices), + completion_tokens, + ) + return ChatCompletionResponse( + id=gen.request_id, + model=request.model, + created=int(time.time()), + choices=choices, + usage=usage, + prompt_logprobs=gen.prompt_logprobs, + kv_transfer_params=gen.kv_transfer_params, + ) + + async def derender_completion_response( + self, + request: DerenderCompletionRequest, + ) -> CompletionResponse | ErrorResponse: + """Postprocess a list of GenerateResponses into a CompletionResponse. + + Non-streaming only. Mirrors the multi-prompt completions case: one + GenerateResponse per prompt, parallel to the list[GenerateRequest] + from /v1/completions/render. + """ + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + return error_check_ret + + ( + choices, + total_prompt_tokens, + total_completion_tokens, + ) = await self.online_derenderer.derender_completion( + request.generate_responses, request.prompt_tokens + ) + + if not request.generate_responses: + return self.create_error_response("generate_responses must not be empty") + + first = request.generate_responses[0] + kv_params = first.kv_transfer_params + if any( + r.kv_transfer_params != kv_params for r in request.generate_responses[1:] + ): + logger.warning( + "derender_completion: kv_transfer_params differ across responses; " + "setting to None on the aggregated response" + ) + kv_params = None + + usage = UsageInfo( + prompt_tokens=total_prompt_tokens, + completion_tokens=total_completion_tokens, + total_tokens=total_prompt_tokens + total_completion_tokens, + ) + + logger.debug( + "derender_completion request_id=%s model=%s choices=%d" + " completion_tokens=%d", + first.request_id, + request.model, + len(choices), + total_completion_tokens, + ) + return CompletionResponse( + id=first.request_id, + model=request.model, + created=int(time.time()), + choices=choices, + usage=usage, + kv_transfer_params=kv_params, + ) + + @staticmethod + def _extract_mm_features( + engine_input: EngineInput, + ) -> MultiModalFeatures | None: + """Extract multimodal metadata from a rendered engine prompt. + + Returns ``None`` for text-only prompts. + """ + if engine_input.get("type") != "multimodal": + return None + + # 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: [ + PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges + ] + 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, + ) diff --git a/vllm/entrypoints/scale_out/factories.py b/vllm/entrypoints/scale_out/factories.py new file mode 100644 index 00000000000..341dad13a86 --- /dev/null +++ b/vllm/entrypoints/scale_out/factories.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace +from typing import TYPE_CHECKING + +from fastapi import FastAPI + +from vllm.engine.protocol import EngineClient +from vllm.tasks import SupportedTask + +if TYPE_CHECKING: + from starlette.datastructures import State + + from vllm.entrypoints.serve.utils.request_logger import RequestLogger +else: + RequestLogger = object + + +def init_render_state( + state: "State", + request_logger: RequestLogger | None, +): + from .derender.serving import ServingDerender + from .render.serving import ServingRender + + state.serving_render = ServingRender( + state.openai_serving_models, + state.online_renderer, + request_logger=request_logger, + ) + + state.serving_derender = ServingDerender( + state.openai_serving_models, + state.online_derenderer, + request_logger=request_logger, + ) + + +def init_scale_out_state( + state: "State", + args: "Namespace", + engine_client: "EngineClient", + request_logger: RequestLogger | None, +): + init_render_state(state, request_logger) + + from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens + + state.serving_tokens = ServingTokens( + engine_client, + state.openai_serving_models, + state.online_renderer, + request_logger=request_logger, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_log_outputs=args.enable_log_outputs, + force_no_detokenize=args.tokens_only, + ) + + +def register_scale_out_api_routers( + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], +): + from .render.api_router import router as render_render + + app.include_router(render_render) + + from .derender.api_router import router as derender_render + + app.include_router(derender_render) + + if "generate" in supported_tasks: + from .token_in_token_out.api_router import ( + attach_router as attach_disagg_router, + ) + + attach_disagg_router(app) diff --git a/vllm/entrypoints/serve/render/__init__.py b/vllm/entrypoints/scale_out/render/__init__.py similarity index 100% rename from vllm/entrypoints/serve/render/__init__.py rename to vllm/entrypoints/scale_out/render/__init__.py diff --git a/vllm/entrypoints/serve/render/api_router.py b/vllm/entrypoints/scale_out/render/api_router.py similarity index 50% rename from vllm/entrypoints/serve/render/api_router.py rename to vllm/entrypoints/scale_out/render/api_router.py index 3b3ad476124..d2452866446 100644 --- a/vllm/entrypoints/serve/render/api_router.py +++ b/vllm/entrypoints/scale_out/render/api_router.py @@ -2,27 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from http import HTTPStatus -from fastapi import APIRouter, Depends, FastAPI, Request +from fastapi import APIRouter, Depends, Request from fastapi.responses import JSONResponse -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionResponse, -) -from vllm.entrypoints.openai.completion.protocol import ( - CompletionRequest, - CompletionResponse, -) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.serve.disagg.protocol import ( - DerenderChatRequest, - DerenderCompletionRequest, - GenerateRequest, -) -from vllm.entrypoints.serve.render.serving import ServingRender from vllm.entrypoints.serve.utils.api_utils import validate_json_request from vllm.logger import init_logger +from ..token_in_token_out.protocol import GenerateRequest +from .serving import ServingRender + logger = init_logger(__name__) router = APIRouter() @@ -79,55 +70,3 @@ async def render_completion(request: CompletionRequest, raw_request: Request): return JSONResponse(content=result.model_dump(), status_code=result.error.code) return JSONResponse(content=[item.model_dump() for item in result]) - - -@router.post( - "/v1/chat/completions/derender", - dependencies=[Depends(validate_json_request)], - response_model=ChatCompletionResponse, - responses={ - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, -) -async def derender_chat_completion(request: DerenderChatRequest, raw_request: Request): - handler = render(raw_request) - if handler is None: - raise NotImplementedError( - "The model does not support Chat Completions Derender API" - ) - - result = await handler.derender_chat_response(request) - - if isinstance(result, ErrorResponse): - return JSONResponse(content=result.model_dump(), status_code=result.error.code) - - return JSONResponse(content=result.model_dump()) - - -@router.post( - "/v1/completions/derender", - dependencies=[Depends(validate_json_request)], - response_model=CompletionResponse, - responses={ - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, -) -async def derender_completion(request: DerenderCompletionRequest, raw_request: Request): - handler = render(raw_request) - if handler is None: - raise NotImplementedError("The model does not support Completions Derender API") - - result = await handler.derender_completion_response(request) - - if isinstance(result, ErrorResponse): - return JSONResponse(content=result.model_dump(), status_code=result.error.code) - - return JSONResponse(content=result.model_dump()) - - -def attach_router(app: FastAPI) -> None: - app.include_router(router) diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py similarity index 66% rename from vllm/entrypoints/serve/render/serving.py rename to vllm/entrypoints/scale_out/render/serving.py index adcaf8af9af..105bd75185d 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -1,28 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import time from typing import cast -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionResponse, -) -from vllm.entrypoints.openai.completion.protocol import ( - CompletionRequest, - CompletionResponse, -) -from vllm.entrypoints.openai.engine.protocol import ( - ErrorResponse, - UsageInfo, -) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.models.serving import ( OpenAIModelRegistry, OpenAIServingModels, ) -from vllm.entrypoints.serve.disagg.mm_serde import encode_mm_kwargs_item -from vllm.entrypoints.serve.disagg.protocol import ( - DerenderChatRequest, - DerenderCompletionRequest, +from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import encode_mm_kwargs_item +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, MultiModalFeatures, PlaceholderRangeInfo, @@ -41,7 +29,6 @@ from vllm.renderers.inputs.preprocess import ( extract_prompt_components, extract_prompt_len, ) -from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.utils import random_uuid @@ -53,7 +40,6 @@ class ServingRender(BaseServing): self, models: OpenAIServingModels | OpenAIModelRegistry, online_renderer: "OnlineRenderer", - online_derenderer: "OnlineDerenderer", *, request_logger: RequestLogger | None = None, ) -> None: @@ -64,7 +50,6 @@ class ServingRender(BaseServing): ) self.online_renderer = online_renderer - self.online_derenderer = online_derenderer self.default_sampling_params = ( online_renderer.model_config.get_diff_sampling_param() @@ -223,117 +208,6 @@ class ServingRender(BaseServing): return generate_requests - async def derender_chat_response( - self, - request: DerenderChatRequest, - ) -> ChatCompletionResponse | ErrorResponse: - """Postprocess a GenerateResponse into a ChatCompletionResponse. - - Non-streaming only: expects the complete GenerateResponse with all - token IDs present. Uses ``parser.parse()`` for one-shot extraction. - - When ``request.chat_request`` is provided, the parser splits the - output into (reasoning, content, tool_calls). Otherwise falls - back to plain detokenization. - """ - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - return error_check_ret - - try: - choices = await self.online_derenderer.derender_chat( - request.generate_response, request.chat_request - ) - except ValueError as exc: - return self.create_error_response(str(exc)) - - prompt_tokens = ( - request.prompt_tokens if request.prompt_tokens is not None else 0 - ) - gen = request.generate_response - completion_tokens = sum(len(ch.token_ids) for ch in gen.choices if ch.token_ids) - usage = UsageInfo( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) - - logger.debug( - "derender_chat request_id=%s model=%s choices=%d completion_tokens=%d", - gen.request_id, - request.model, - len(choices), - completion_tokens, - ) - return ChatCompletionResponse( - id=gen.request_id, - model=request.model, - created=int(time.time()), - choices=choices, - usage=usage, - prompt_logprobs=gen.prompt_logprobs, - kv_transfer_params=gen.kv_transfer_params, - ) - - async def derender_completion_response( - self, - request: DerenderCompletionRequest, - ) -> CompletionResponse | ErrorResponse: - """Postprocess a list of GenerateResponses into a CompletionResponse. - - Non-streaming only. Mirrors the multi-prompt completions case: one - GenerateResponse per prompt, parallel to the list[GenerateRequest] - from /v1/completions/render. - """ - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - return error_check_ret - - ( - choices, - total_prompt_tokens, - total_completion_tokens, - ) = await self.online_derenderer.derender_completion( - request.generate_responses, request.prompt_tokens - ) - - if not request.generate_responses: - return self.create_error_response("generate_responses must not be empty") - - first = request.generate_responses[0] - kv_params = first.kv_transfer_params - if any( - r.kv_transfer_params != kv_params for r in request.generate_responses[1:] - ): - logger.warning( - "derender_completion: kv_transfer_params differ across responses; " - "setting to None on the aggregated response" - ) - kv_params = None - - usage = UsageInfo( - prompt_tokens=total_prompt_tokens, - completion_tokens=total_completion_tokens, - total_tokens=total_prompt_tokens + total_completion_tokens, - ) - - logger.debug( - "derender_completion request_id=%s model=%s choices=%d" - " completion_tokens=%d", - first.request_id, - request.model, - len(choices), - total_completion_tokens, - ) - return CompletionResponse( - id=first.request_id, - model=request.model, - created=int(time.time()), - choices=choices, - usage=usage, - kv_transfer_params=kv_params, - ) - @staticmethod def _extract_mm_features( engine_input: EngineInput, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/__init__.py b/vllm/entrypoints/scale_out/token_in_token_out/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/serve/disagg/api_router.py b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py similarity index 96% rename from vllm/entrypoints/serve/disagg/api_router.py rename to vllm/entrypoints/scale_out/token_in_token_out/api_router.py index e5bd351e01f..30857e4c1cf 100644 --- a/vllm/entrypoints/serve/disagg/api_router.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py @@ -11,13 +11,6 @@ from fastapi.responses import JSONResponse, StreamingResponse from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.serve.disagg.protocol import ( - GenerateRequest, - GenerateResponse, -) -from vllm.entrypoints.serve.disagg.serving import ( - ServingTokens, -) from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -26,6 +19,12 @@ from vllm.entrypoints.serve.utils.api_utils import ( ) from vllm.logger import init_logger +from .protocol import ( + GenerateRequest, + GenerateResponse, +) +from .serving import ServingTokens + logger = init_logger(__name__) diff --git a/vllm/entrypoints/serve/disagg/mm_serde.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py similarity index 100% rename from vllm/entrypoints/serve/disagg/mm_serde.py rename to vllm/entrypoints/scale_out/token_in_token_out/mm_serde.py diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py similarity index 100% rename from vllm/entrypoints/serve/disagg/protocol.py rename to vllm/entrypoints/scale_out/token_in_token_out/protocol.py diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py similarity index 99% rename from vllm/entrypoints/serve/disagg/serving.py rename to vllm/entrypoints/scale_out/token_in_token_out/serving.py index cbd6f83f233..70185a85b30 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -28,14 +28,6 @@ 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, - GenerateResponseChoice, - GenerateResponseStreamChoice, - GenerateStreamResponse, -) from vllm.entrypoints.serve.utils.api_utils import get_max_tokens, should_include_usage from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.inputs import EngineInput, mm_input @@ -51,6 +43,15 @@ from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils.collection_utils import as_list +from .mm_serde import decode_mm_kwargs_item +from .protocol import ( + GenerateRequest, + GenerateResponse, + GenerateResponseChoice, + GenerateResponseStreamChoice, + GenerateStreamResponse, +) + logger = init_logger(__name__) diff --git a/vllm/entrypoints/serve/engine/typing.py b/vllm/entrypoints/serve/engine/typing.py index 8f0b7835dab..2e01c092c7b 100644 --- a/vllm/entrypoints/serve/engine/typing.py +++ b/vllm/entrypoints/serve/engine/typing.py @@ -15,7 +15,7 @@ from vllm.entrypoints.openai.completion.protocol import ( CompletionResponse, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.entrypoints.serve.disagg.protocol import ( +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderChatRequest, DerenderCompletionRequest, GenerateRequest, diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 91d03bbe819..fb4b880c48c 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -16,7 +16,7 @@ from vllm.entrypoints.openai.completion.protocol import ( ) from vllm.entrypoints.openai.engine.protocol import ToolCall from vllm.entrypoints.openai.engine.serving import resolve_token_id_placeholder -from vllm.entrypoints.serve.disagg.protocol import GenerateResponse +from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateResponse from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger from vllm.parser import Parser, ParserManager From 59575da46df964e6161fb0e1a77fa76ea9ce3106 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 29 Jun 2026 20:30:28 +0800 Subject: [PATCH 114/138] [XPU] exclude unsupported models for test_tensor_sechma.py (#47008) Signed-off-by: Yan Ma Signed-off-by: Kunshang Ji Co-authored-by: Kunshang Ji --- .../intel_jobs/models_multimodal_intel.yaml | 2 -- .../models/multimodal/processing/test_common.py | 16 +++++++++++++++- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 0e126906044..42d429f007f 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -125,7 +125,5 @@ steps: pip install open-clip-torch --no-deps && cd tests && pytest -v -s models/multimodal/processing/test_tensor_schema.py - --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]" - --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB' parallelism: 4 diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index ea5aeb8e2ca..f785a68f977 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -20,6 +20,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import MultiModalProcessorOnlyCache from vllm.multimodal.inputs import batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext +from vllm.platforms import current_platform from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config from vllm.utils.mistral import is_mistral_tokenizer @@ -83,6 +84,12 @@ MM_DATA_PATCHES = { "glmasr": glmasr_patch_mm_data, } +_XPU_EXCLUDED_MODEL_IDS = { + "baidu/Unlimited-OCR", + "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4", + "Qwen/Qwen2.5-Omni-7B-AWQ", +} + def _iter_model_ids_to_test(model_arch_list: AbstractSet[str]): for model_arch in model_arch_list: @@ -97,7 +104,14 @@ def _iter_model_ids_to_test(model_arch_list: AbstractSet[str]): def _get_model_ids_to_test(model_arch_list: AbstractSet[str]): - return list(_iter_model_ids_to_test(model_arch_list)) + model_ids = list(_iter_model_ids_to_test(model_arch_list)) + + if current_platform.is_xpu(): + for excluded_model_id in _XPU_EXCLUDED_MODEL_IDS: + while excluded_model_id in model_ids: + model_ids.remove(excluded_model_id) + + return model_ids def get_model_ids_to_test(): From bc8481af09cd4c7f7272ba7bc1913f1051649813 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Mon, 29 Jun 2026 09:19:29 -0400 Subject: [PATCH 115/138] [MoE Refactor] Standardize Humming MoE experts + utilities (#43373) Signed-off-by: Bill Nell Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../fused_moe/experts/fused_humming_moe.py | 325 ++++-- .../layers/fused_moe/oracle/mxfp4.py | 17 +- .../compressed_tensors_moe_w4a4_mxfp4.py | 1 + .../model_executor/layers/quantization/fp8.py | 18 +- .../layers/quantization/humming.py | 234 ++--- .../layers/quantization/quark/quark_moe.py | 1 + .../quantization/utils/humming_utils.py | 957 ++++++++++++++++-- vllm/utils/humming.py | 13 + 8 files changed, 1176 insertions(+), 390 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 5177fa0cde4..047ae46c0d3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -36,28 +36,42 @@ from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, swiglu_limit_func, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8DynamicTokenSym, + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kInt4Static, + kInt8Static, + kMxfp4Dynamic, + kMxfp4Static, + kMxfp8Dynamic, + kMxfp8Static, + kNvfp4Static, +) from vllm.platforms import current_platform -from vllm.utils.humming import GemmType as HummingGemmType -from vllm.utils.humming import HummingLayerMeta, HummingMethod, dtypes +from vllm.utils.import_utils import has_humming from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.utils.humming import GemmType as HummingGemmType logger = init_logger(__name__) -def get_humming_moe_gemm_type() -> str: - env_gemm_type: str = envs.VLLM_HUMMING_MOE_GEMM_TYPE or "" - env_gemm_type = env_gemm_type.lower() - if env_gemm_type == "indexed": - gemm_type = env_gemm_type - elif env_gemm_type in ["grouped_contiguous", "grouped"]: - gemm_type = "grouped_contiguous" - else: - gemm_type = "indexed" +def get_humming_moe_gemm_type() -> str | None: + env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE + gemm_type = None + if env_gemm_type is not None: + env_gemm_type = env_gemm_type.lower() + if env_gemm_type == "indexed": + gemm_type = env_gemm_type + elif env_gemm_type in ["grouped_contiguous", "grouped"]: + gemm_type = "grouped_contiguous" + else: + gemm_type = "indexed" logger.info_once(f"Using {gemm_type} gemm for humming moe") # noqa return gemm_type @@ -89,6 +103,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): self._permute_scratch: MoEPermuteScratch | None = None def init_humming_moe(self): + from vllm.utils.humming import HummingMethod + self.compute_config = { "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, @@ -141,7 +157,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): return math.ceil(global_valid_shape_m * num_experts / global_num_experts) @staticmethod - def humming_gemm_type() -> HummingGemmType: + def humming_gemm_type() -> "HummingGemmType": raise NotImplementedError @classmethod @@ -153,12 +169,51 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + (kMxfp4Static, kMxfp4Dynamic), + (kMxfp4Static, kMxfp8Dynamic), + (kMxfp4Static, kFp8DynamicTokenSym), + (kNvfp4Static, None), + (kNvfp4Static, kFp8DynamicTokenSym), + (kMxfp8Static, None), + (kMxfp8Static, kFp8DynamicTokenSym), + (kFp8StaticChannelSym, None), + (kFp8StaticChannelSym, kFp8DynamicTokenSym), + (kFp8Static128BlockSym, None), + (kFp8Static128BlockSym, kFp8DynamicTokenSym), + (kInt4Static, None), + (kInt4Static, kFp8DynamicTokenSym), + (kInt8Static, None), + (kInt8Static, kFp8DynamicTokenSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @property + def expects_unquantized_inputs(self) -> bool: + """ + Humming kernels handle input quantization internally via + HummingMethod.may_quant_input() in the apply() method. + + This property tells the prepare/finalize step to skip input + quantization (by setting defer_input_quant=True) and pass + unquantized inputs to the experts. This prevents double + quantization: once in prepare and once in Humming's apply(). + + Returns: + True to indicate that this expert expects unquantized inputs + and will handle quantization internally. + """ return True @staticmethod def _supports_current_device() -> bool: platform = current_platform - return platform.is_cuda() and platform.has_device_capability((7, 5)) + return ( + has_humming() + and platform.is_cuda() + and platform.has_device_capability((7, 5)) + ) @staticmethod def _supports_no_act_and_mul() -> bool: @@ -182,10 +237,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return not ( - moe_parallel_config.use_fi_nvl_two_sided_kernels - or moe_parallel_config.use_fi_nvl_one_sided_kernels - ) + return True def moe_problem_size( self, @@ -194,6 +246,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): w2: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple[int, int, int, int, int]: + from vllm.utils.humming import HummingLayerMeta + meta1: HummingLayerMeta = self.layer.humming_metas["w13"] meta2: HummingLayerMeta = self.layer.humming_metas["w2"] @@ -215,6 +269,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): + from vllm.utils.humming import GemmType as HummingGemmType + from vllm.utils.humming import dtypes + num_experts = self.num_experts N = self.layer.intermediate_size_per_partition K = self.layer.hidden_size @@ -254,7 +311,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): torch_dtype_map = { dtypes.float16: torch.float16, dtypes.bfloat16: torch.bfloat16, + dtypes.float32: torch.float32, dtypes.float8e4m3: torch.float8_e4m3fn, + dtypes.float8e5m2: torch.float8_e5m2, dtypes.int8: torch.int8, dtypes.int4: torch.uint8, } @@ -289,7 +348,13 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): for key in buffer_metas: meta = buffer_metas[key] if "quanted" in key and a_dtype.num_bits == 4: - meta["shape"] = meta["shape"][:-1] + (meta["shape"][-1] // 2,) + last_dim = meta["shape"][-1] + if last_dim % 2 != 0: + raise ValueError( + f"Int4 packing requires last dimension to be even, " + f"got {last_dim} for buffer '{key}'" + ) + meta["shape"] = meta["shape"][:-1] + (last_dim // 2,) if num_bits == 16: required_buffers = ["gate_up_output", "activation_output", "down_output"] @@ -325,8 +390,13 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): output_key = "down_output" if self.is_batched() else "output" output_shape = buffer_metas[output_key]["shape"] + elem_size = self.layer.params_dtype.itemsize - return (workspace1_nbytes // 2,), (workspace2_nbytes // 2,), output_shape + return ( + (workspace1_nbytes // elem_size,), + (workspace2_nbytes // elem_size,), + output_shape, + ) def workspace_shapes( self, @@ -344,7 +414,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): def make_workspaces(self, M: int, topk: int, activation: MoEActivation): shapes = self._workspace_shapes(M, topk, activation) workspace1_shape, workspace2_shape, output_shape = shapes - torch_dtype = self.layer.param_dtype + torch_dtype = self.layer.params_dtype workspace1, workspace2 = current_workspace_manager().get_simultaneous( (workspace1_shape, torch_dtype), (workspace2_shape, torch_dtype), @@ -370,45 +440,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): return buffers - 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, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, - ): - assert not apply_router_weight_on_input - - self.main_apply( - hidden_states=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - workspace1=workspace13, - workspace2=workspace2, - expert_tokens_meta=expert_tokens_meta, - ) - - def main_apply( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - workspace1: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - ): - raise NotImplementedError + # Note: apply method is implemented by subclasses following the + # standard FusedMoEExpertsModular.apply signature @staticmethod def is_supported_config( @@ -418,24 +451,27 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, ) -> tuple[bool, str | None]: - if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: - supported = cls.activation_format() == activation_format - reason = "activation_format mismatched" - elif activation_format == mk.FusedMoEActivationFormat.Standard: - if cls.activation_format() != mk.FusedMoEActivationFormat.Standard: - supported = False - reason = "activation_format mismatched" - else: - assert hasattr(cls, "humming_gemm_type") - gemm_type = cls.humming_gemm_type().value.lower() - preferred_gemm_type = get_humming_moe_gemm_type().lower() - supported = preferred_gemm_type == gemm_type - reason = "preferred gemm type mismatched" - else: - supported = False - reason = "unsupported activation_format" + supported, reason = mk.FusedMoEExpertsModular.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) - return supported, None if supported else reason + if supported: + assert hasattr(cls, "humming_gemm_type") + gemm_type = cls.humming_gemm_type().value.lower() + preferred_gemm_type = get_humming_moe_gemm_type() + if preferred_gemm_type is not None: + supported = preferred_gemm_type.lower() == gemm_type + if not supported: + reason = ( + f"preferred gemm type {preferred_gemm_type} != " + f"supported gemm type {gemm_type}" + ) + + return supported, reason def apply_activation( self, @@ -459,7 +495,9 @@ class HummingIndexedExperts(HummingExpertsBase): return mk.FusedMoEActivationFormat.Standard @staticmethod - def humming_gemm_type() -> HummingGemmType: + def humming_gemm_type() -> "HummingGemmType": + from vllm.utils.humming import GemmType as HummingGemmType + return HummingGemmType.INDEXED def prepare_humming_moe_kwargs( @@ -470,12 +508,18 @@ class HummingIndexedExperts(HummingExpertsBase): ) -> tuple[dict[str, Any], dict[str, Any]]: valid_shape_m = self.estimate_local_valid_shape_m(topk_ids) + moe_block_size = None for min_shape_m, max_shape_m, config in self.w13_tuning_config: if valid_shape_m > min_shape_m and valid_shape_m <= max_shape_m: moe_block_size = config["block_shape"][0] break - else: - raise ValueError(f"cannot found moe_block_size for shape {valid_shape_m}") + + if moe_block_size is None: + logger.warning_once( + "No tuning config found for shape %s, using default block_size=64", + valid_shape_m, + ) + moe_block_size = 64 sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size( topk_ids=topk_ids, @@ -501,27 +545,47 @@ class HummingIndexedExperts(HummingExpertsBase): return moe_kwargs1, moe_kwargs2 - def main_apply( + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - workspace1: 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, workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, - ): + apply_router_weight_on_input: bool, + ) -> None: + """ + Standard apply implementation for Humming indexed experts. + + Note: Humming kernels handle weights and quantization internally through + the layer object, so w1, w2, a1q_scale, a2_scale parameters are not used. + The output is written into workspace13 via the buffer management. + """ + from vllm.utils.humming import HummingMethod + + assert not apply_router_weight_on_input + hidden_states = hidden_states.view(-1, hidden_states.size(-1)) buffers = self.prepare_buffers( - workspace1, + workspace13, workspace2, topk_ids.size(0), topk_ids.size(1), - self.layer.activation, + activation, ) moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs( topk_ids=topk_ids, - expert_map=self.layer.expert_map, + expert_map=expert_map, expert_tokens_meta=expert_tokens_meta, ) @@ -542,7 +606,7 @@ class HummingIndexedExperts(HummingExpertsBase): ) self.apply_activation( - activation=self.layer.activation, + activation=activation, input=buffers["gate_up_output"], output=buffers["activation_output"], ) @@ -567,10 +631,13 @@ class HummingIndexedExperts(HummingExpertsBase): inputs=buffers["down_output"].view(*topk_ids.shape, -1), topk_weights=topk_weights, topk_ids=topk_ids, - expert_map=self.layer.expert_map, + expert_map=expert_map, outputs=buffers["output"], ) + # Note: output is already written to buffers["output"] + # which aliases workspace13/output + class HummingGroupedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -581,35 +648,57 @@ class HummingGroupedExperts(HummingExpertsBase): return mk.FusedMoEActivationFormat.Standard @staticmethod - def humming_gemm_type() -> HummingGemmType: + def humming_gemm_type() -> "HummingGemmType": + from vllm.utils.humming import GemmType as HummingGemmType + return HummingGemmType.GROUPED_CONTIGUOUS - def main_apply( + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - workspace1: 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, workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, - ): + apply_router_weight_on_input: bool, + ) -> None: + """ + Standard apply implementation for Humming grouped experts. + + Note: Humming kernels handle weights and quantization internally through + the layer object, so w1, w2, a1q_scale, a2_scale parameters are not used. + The output is written into workspace13 via the buffer management. + """ + from vllm.utils.humming import HummingMethod + + assert not apply_router_weight_on_input + valid_shape_m = self.estimate_local_valid_shape_m(topk_ids) buffers = self.prepare_buffers( - workspace1, + workspace13, workspace2, topk_ids.size(0), topk_ids.size(1), - self.layer.activation, + activation, ) hidden_states, _, expert_first_token_offset, inv_perm, _ = moe_permute( hidden_states=hidden_states, a1q_scale=None, topk_ids=topk_ids, - n_expert=self.global_num_experts, + n_expert=global_num_experts, n_local_expert=self.num_experts, - expert_map=self.layer.expert_map, + expert_map=expert_map, scratch=self._get_permute_scratch(), ) @@ -633,7 +722,7 @@ class HummingGroupedExperts(HummingExpertsBase): ) self.apply_activation( - activation=self.layer.activation, + activation=activation, input=buffers["gate_up_output"], output=buffers["activation_output"], ) @@ -665,6 +754,9 @@ class HummingGroupedExperts(HummingExpertsBase): expert_first_token_offset=expert_first_token_offset, ) + # Note: output is already written to buffers["output"] + # which aliases workspace13/output + class BatchedHummingGroupedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -675,29 +767,51 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): return mk.FusedMoEActivationFormat.BatchedExperts @staticmethod - def humming_gemm_type() -> HummingGemmType: + def humming_gemm_type() -> "HummingGemmType": + from vllm.utils.humming import GemmType as HummingGemmType + return HummingGemmType.GROUPED_MASKED - def main_apply( + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - workspace1: 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, workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, - ): + apply_router_weight_on_input: bool, + ) -> None: + """ + Standard apply implementation for Humming batched grouped experts. + + Note: Humming kernels handle weights and quantization internally through + the layer object, so w1, w2, a1q_scale, a2_scale parameters are not used. + The output is written into workspace13 via the buffer management. + """ + from vllm.utils.humming import HummingMethod + + assert not apply_router_weight_on_input assert expert_tokens_meta is not None + hidden_states = hidden_states.view(-1, hidden_states.size(-1)) valid_shape_m = self.estimate_local_valid_shape_m(topk_ids) expert_num_tokens = expert_tokens_meta.expert_num_tokens buffers = self.prepare_buffers( - workspace1, + workspace13, workspace2, topk_ids.size(0), topk_ids.size(1), - self.layer.activation, + activation, ) inputs, input_scale = HummingMethod.may_quant_input( @@ -720,7 +834,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): ) self.apply_activation( - activation=self.layer.activation, + activation=activation, input=buffers["gate_up_output"], output=buffers["activation_output"], ) @@ -743,3 +857,6 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): tuning_config=self.w2_tuning_config_str, sublayer_name="w2", ) + + # Note: output is already written to buffers["down_output"] + # which aliases workspace13/output diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index b1b41ded11a..55a767f060c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -711,10 +711,12 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( if mxfp4_backend == Mxfp4MoeBackend.HUMMING: from vllm.model_executor.layers.quantization.utils.humming_utils import ( - prepare_humming_moe_layer, + convert_to_humming_moe_kernel_format, ) - prepare_humming_moe_layer(layer, {"quant_method": "gpt_oss_mxfp4"}) + convert_to_humming_moe_kernel_format( + layer, quant_config={"quant_method": "gpt_oss_mxfp4"} + ) return ( layer.w13_weight, layer.w2_weight, @@ -1277,10 +1279,12 @@ def convert_weight_to_mxfp4_moe_kernel_format( if mxfp4_backend == Mxfp4MoeBackend.HUMMING: from vllm.model_executor.layers.quantization.utils.humming_utils import ( - prepare_humming_moe_layer, + convert_to_humming_moe_kernel_format, ) - prepare_humming_moe_layer(layer, {"quant_method": "mxfp4"}) + convert_to_humming_moe_kernel_format( + layer, quant_config={"quant_method": "mxfp4"} + ) return ( layer.w13_weight, layer.w2_weight, @@ -1569,7 +1573,7 @@ def make_mxfp4_moe_quant_config( w2_bias: torch.Tensor | None = None, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, - layer: torch.nn.Module | None = None, + layer: "RoutedExperts | None" = None, ) -> FusedMoEQuantConfig | None: """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: @@ -1662,7 +1666,7 @@ def make_mxfp4_moe_quant_config( get_humming_moe_quant_config, ) - assert isinstance(layer, RoutedExperts) + assert layer is not None return get_humming_moe_quant_config( layer, gemm1_alpha=gemm1_alpha, @@ -1703,6 +1707,7 @@ def make_mxfp4_moe_kernel( assert prepare_finalize is not None logger.info_once("Using %s", prepare_finalize.__class__.__name__) + logger.info_once("Using %s", experts_cls.__name__) extra_kwargs = {} if mxfp4_backend == Mxfp4MoeBackend.HUMMING: 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 906b0727b18..1a0cd8ed100 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 @@ -143,6 +143,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): mxfp4_backend=self.mxfp4_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, + layer=layer, ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index a71754769a4..86818ed4b7e 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -706,15 +706,15 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer.w2_weight.is_shuffled = True 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._expert_routing_tables(), - ) + assert self.moe_quant_config is not None + 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._expert_routing_tables(), + ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: # Allow for accessing weights and scales in standard way. diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index eb598ff7f79..9bbbf41115e 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -30,6 +30,14 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + get_humming_moe_quant_config, + input_schema_to_quant_key, + make_humming_moe_kernel, + select_humming_moe_experts, + weight_schema_to_quant_key, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.parameter import ( BasevLLMParameter, @@ -106,7 +114,7 @@ def prepare_param(tensor, name, extra_attrs): return param -def prepare_moe_param(tensor, name, extra_attrs): +def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, Any]): param = torch.nn.Parameter(tensor, requires_grad=False) if "scale_type" in extra_attrs: extra_attrs["quant_method"] = extra_attrs["scale_type"] @@ -605,12 +613,27 @@ class HummingMoEMethod(FusedMoEMethodBase): ) -> None: super().__init__(moe) self.quant_config = quant_config - self.moe = moe self.weight_schema = quant_config.weight_schema self.input_schema = quant_config.input_schema self.force_weight_schema = quant_config.force_weight_schema self.force_input_schema = quant_config.force_input_schema + # Derive QuantKeys from humming schemas. + # Prefer force schemas (the final format after requant) over base. + weight_key = weight_schema_to_quant_key( + self.force_weight_schema or self.weight_schema + ) + activation_key = input_schema_to_quant_key( + self.force_input_schema or self.input_schema + ) + + # Select Humming MoE experts + self.experts_cls = select_humming_moe_experts( + config=self.moe, + weight_key=weight_key, + activation_key=activation_key, + ) + def prepare_weight_loader(self, layer, weight_loader): def new_weight_loader( param: torch.nn.Parameter, @@ -647,7 +670,7 @@ class HummingMoEMethod(FusedMoEMethodBase): sublayer_name = "w2" if shard_id == "w2" else "w13" param = getattr(layer, sublayer_name + "_" + key) - part_subccess = param.weight_loader( + part_success = param.weight_loader( param=param, loaded_weight=tensor.cpu(), weight_name=shard_id + "_" + key, @@ -655,7 +678,7 @@ class HummingMoEMethod(FusedMoEMethodBase): expert_id=expert_id, return_success=return_success, ) - success = success and part_subccess + success = success and part_success return success if return_success else None @@ -677,7 +700,7 @@ class HummingMoEMethod(FusedMoEMethodBase): def create_weights( self, - layer: torch.nn.Module, + layer: RoutedExperts, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, @@ -734,160 +757,35 @@ class HummingMoEMethod(FusedMoEMethodBase): locks = torch.zeros(1024, dtype=torch.int32) layer.register_buffer("locks", locks) - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - get_humming_moe_quant_config, - ) - + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: return get_humming_moe_quant_config(layer) def process_weights_after_loading(self, layer: RoutedExperts) -> None: if getattr(self, "processed", False): return self.processed = True - layer.weight_schemas = {} - layer.input_schemas = {} - for sublayer_name, configs in layer.sublayer_configs.items(): - input_schema = self.input_schema - weight_schema = self.weight_schema - # convert from checkpoint format to humming format - if not isinstance(weight_schema, _hm.HummingWeightSchema): - tensors: dict[str, torch.Tensor] = dict( - (key.removeprefix(sublayer_name + "_"), value) - for key, value in layer.state_dict().items() - if key.startswith(sublayer_name + "_") - ) - shape_k_stacks = [configs["shape_k"]] - shape_n_stacks = [configs["shape_n"]] - if sublayer_name == "w13": - shape_n_stacks = [configs["shape_n"] // 2] * 2 - - weight_schema, tensors = weight_schema.convert_humming( - tensors=tensors, - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - param_dtype=layer.param_dtype, - num_experts=layer.num_experts, - ) - - input_schema, _ = input_schema.convert_humming( - tensors=tensors, - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - param_dtype=layer.param_dtype, - num_experts=layer.num_experts, - ) - - for name, _ in list(layer.named_parameters()): - if not name.startswith(sublayer_name + "_"): - continue - delattr(layer, name) - - for name, tensor in tensors.items(): - name = f"{sublayer_name}_{name}" - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - - layer.weight_schemas[sublayer_name] = weight_schema - layer.input_schemas[sublayer_name] = input_schema - - # force requant (origin quant setting -> fp16/bf16 -> new_quant setting) - assert isinstance(weight_schema, _hm.HummingWeightSchema) - force_requant = self.force_weight_schema is not None - if force_requant and weight_schema != self.force_weight_schema: - tensors = dict( - (key.removeprefix(sublayer_name + "_"), value) - for key, value in layer.state_dict().items() - if key.startswith(sublayer_name + "_") - ) - - tensors = weight_schema.requant_tensors( - tensors=tensors, - target_weight_schema=self.force_weight_schema, - param_dtype=layer.param_dtype, - ) - - weight_schema = self.force_weight_schema - - for name, _ in list(layer.named_parameters()): - if not name.startswith(sublayer_name + "_"): - continue - if name == sublayer_name + "_bias": - continue - delattr(layer, name) - - for name, tensor in tensors.items(): - name = f"{sublayer_name}_{name}" - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - - del tensors - - # prepare layer config from humming kernel - _hm.HummingMethod.prepare_layer_meta( - layer=layer, - shape_n=configs["shape_n"], - shape_k=configs["shape_k"], - pad_n_to_multiple=256, - pad_k_to_multiple=128, - input_schema=input_schema, - weight_schema=weight_schema, - has_bias=self.moe.has_bias, - num_experts=layer.num_experts, - torch_dtype=layer.param_dtype, - sublayer_name=sublayer_name, - ) - - # preprocess weight for inference - _hm.HummingMethod.transform_humming_layer( - layer, sublayer_name=sublayer_name - ) - - from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( - HummingGroupedExperts, - HummingIndexedExperts, - get_humming_moe_gemm_type, + # Convert weights to Humming kernel format + convert_to_humming_moe_kernel_format( + layer=layer, + sublayer_configs=layer.sublayer_configs, + weight_schema=self.weight_schema, + input_schema=self.input_schema, + force_weight_schema=self.force_weight_schema, ) - # use moe modular - experts: HummingIndexedExperts | HummingGroupedExperts - layer._ensure_moe_quant_config_init() + # Build the MoE kernel + self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None - if get_humming_moe_gemm_type() == "indexed": - experts = HummingIndexedExperts(layer, self.moe, self.moe_quant_config) - else: - experts = HummingGroupedExperts(layer, self.moe, self.moe_quant_config) - self.experts = experts - - def select_gemm_impl( - self, - prepare_finalize, - layer: torch.nn.Module, - ): - from vllm.model_executor.layers.fused_moe import modular_kernel as mk - from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( - BatchedHummingGroupedExperts, - HummingGroupedExperts, - HummingIndexedExperts, - get_humming_moe_gemm_type, + assert self.experts_cls is not None + self.moe_kernel = make_humming_moe_kernel( + self.moe_quant_config, + self.moe, + self.experts_cls, + layer=layer, + routing_tables=layer._expert_routing_tables(), ) - activation_format = prepare_finalize.activation_format - assert self.moe_quant_config is not None - if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: - return BatchedHummingGroupedExperts( - layer=layer, - moe_config=self.moe, - quant_config=self.moe_quant_config, - max_num_tokens=prepare_finalize.max_num_tokens_per_rank(), - num_dispatchers=prepare_finalize.num_dispatchers(), - ) - elif get_humming_moe_gemm_type() == "indexed": - return HummingIndexedExperts(layer, self.moe, self.moe_quant_config) - else: - return HummingGroupedExperts(layer, self.moe, self.moe_quant_config) - def apply( self, layer: RoutedExperts, @@ -896,22 +794,36 @@ class HummingMoEMethod(FusedMoEMethodBase): topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - workspace1, workspace2, output = self.experts.make_workspaces( - M=topk_ids.size(0), - topk=topk_ids.size(1), - activation=layer.activation, - ) + ) -> torch.Tensor: + """ + Apply Humming-quantized MoE computation using the standard kernel flow. - assert workspace1.data_ptr() == output.data_ptr() + This method uses FusedMoEKernel.apply() which orchestrates: + 1. Preparation (quantization if needed - skipped for Humming via + expects_unquantized_inputs=True to prevent double quantization) + 2. Expert computation (via experts.apply()) + 3. Finalization (weight application & reduction - no-op for Humming + since it's already done internally) - self.experts.main_apply( + Humming handles all quantization, weight application, and reduction + internally in the experts.apply() method via HummingMethod calls. + + Note: Although w1/w2 weights are passed to the kernel for interface + consistency, Humming's experts.apply() reads weights directly from + the layer object via HummingMethod.forward_layer() and ignores the + w1/w2 parameters. + """ + assert self.moe_kernel is not None + return self.moe_kernel.apply( hidden_states=x, - topk_weights=topk_weights, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_ids=topk_ids, - workspace1=workspace1, - workspace2=workspace2, - expert_tokens_meta=None, + topk_weights=topk_weights, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=False, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, ) - - return output diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 70b1e25959e..ebe37d8dc5b 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1280,6 +1280,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) # Emulation and other schemes diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index 617158ae139..d84a2e12f54 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -1,20 +1,376 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json -from typing import Any +from typing import TYPE_CHECKING, Any import regex as re import torch +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import envs -from vllm.model_executor.layers.fused_moe import RoutedExperts +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, FusedMoEQuantDesc, ) +from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, +) from vllm.model_executor.layers.linear import LinearBase -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape -from vllm.utils.humming import BaseWeightSchema, HummingInputSchema, HummingMethod +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP4_DTYPE, + FP8_DTYPE, + INT4_DTYPE, + INT8_DTYPE, + MXFP_SCALE_DTYPE, + GroupShape, + QuantKey, + ScaleDesc, +) +from vllm.utils.import_utils import has_humming + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.utils.humming import ( + AWQWeightSchema, + BaseInputSchema, + BaseWeightSchema, + CompressedTensorsInputSchema, + CompressedTensorsWeightSchema, + Fp8WeightSchema, + GPTQWeightSchema, + HummingInputSchema, + HummingWeightSchema, + ) + from vllm.utils.humming import dtypes as humming_dtypes + +logger = init_logger(__name__) + +if has_humming(): + from vllm.utils.humming import dtypes as humming_dtypes + + _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = { + humming_dtypes.float4e2m1: FP4_DTYPE, + humming_dtypes.float8e4m3: FP8_DTYPE, + humming_dtypes.float8e5m2: torch.float8_e5m2, + humming_dtypes.int8: torch.int8, + humming_dtypes.uint4: INT4_DTYPE, + humming_dtypes.uint8: INT8_DTYPE, + humming_dtypes.uint2: torch.uint8, + humming_dtypes.uint3: torch.uint8, + } + + _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { + humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE, + humming_dtypes.float8e4m3: FP8_DTYPE, + humming_dtypes.float16: torch.float16, + humming_dtypes.bfloat16: torch.bfloat16, + humming_dtypes.float32: torch.float32, + } + + +def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape: + """ + Map humming group sizes to QuantKey GroupShape. + + group_size: elements per group along K (col); 0 means full dimension. + group_size_n: elements per group along N (row); 0 means 1 (per-row). + + GroupShape convention: row = N dim, col = K dim. + """ + if group_size == 0 and group_size_n == 0: + return GroupShape.PER_CHANNEL + + row = group_size_n if group_size_n > 0 else 1 + col = group_size if group_size > 0 else -1 + return GroupShape(row=row, col=col) + + +# ---- HummingWeightSchema (post-conversion) -------------------------------- + + +def _humming_weight_schema_to_quant_key( + schema: "HummingWeightSchema", +) -> QuantKey: + from vllm.utils.humming import WeightScaleType + + """Convert a HummingWeightSchema to a QuantKey.""" + dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype] + + if schema.bs_dtype is not None: + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype] + else: + scale_dtype = torch.float32 + + group_shape = _group_shape( + schema.weight_scale_group_size, + schema.weight_scale_group_size_n, + ) + + scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) + + scale2 = None + if schema.weight_scale_type == WeightScaleType.GROUP_TENSOR: + scale2 = ScaleDesc( + dtype=torch.float32, + static=True, + group_shape=GroupShape.PER_TENSOR, + ) + + return QuantKey( + dtype=dtype, + scale=scale, + scale2=scale2, + symmetric=not schema.has_zero_point, + ) + + +# ---- Checkpoint-format weight schemas (pre-conversion) -------------------- + + +def _fp8_weight_schema_to_quant_key(schema: "Fp8WeightSchema") -> QuantKey: + if schema.weight_block_size is not None: + gs_n, gs_k = schema.weight_block_size + group_shape = GroupShape(row=gs_n, col=gs_k) + else: + group_shape = GroupShape.PER_CHANNEL + + scale = ScaleDesc(dtype=torch.float32, static=True, group_shape=group_shape) + return QuantKey(dtype=FP8_DTYPE, scale=scale, symmetric=True) + + +def _awq_weight_schema_to_quant_key(schema: "AWQWeightSchema") -> QuantKey: + group_shape = _group_shape(schema.group_size) + scale = ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=group_shape, + ) + return QuantKey( + dtype=INT4_DTYPE, + scale=scale, + symmetric=not schema.zero_point, + ) + + +def _gptq_weight_schema_to_quant_key(schema: "GPTQWeightSchema") -> QuantKey: + group_shape = _group_shape(schema.group_size) + scale = ScaleDesc( + dtype=torch.float16, + static=True, + group_shape=group_shape, + ) + return QuantKey(dtype=INT4_DTYPE, scale=scale, symmetric=schema.sym) + + +def _compressed_tensors_weight_schema_to_quant_key( + schema: "CompressedTensorsWeightSchema", +) -> QuantKey: + # Determine dtype from format/type/num_bits + fmt = schema.format + if fmt in ("int-quantized", "float-quantized", "naive-quantized"): + dtype = INT8_DTYPE if schema.type == "int" else FP8_DTYPE + elif "nvfp4" in fmt or "mxfp4" in fmt: + dtype = FP4_DTYPE + else: + dtype = _HUMMING_TO_QUANT_DTYPE[ + humming_dtypes.DataType.from_str(f"uint{schema.num_bits}") + ] + + # Determine group shape from strategy + if schema.strategy in ("group", "tensor_group"): + group_shape = _group_shape(schema.group_size or 0) + elif schema.strategy == "block" and schema.block_structure is not None: + group_shape = GroupShape( + row=schema.block_structure[0], + col=schema.block_structure[1], + ) + else: + group_shape = GroupShape.PER_CHANNEL + + # Determine scale dtype + if "mxfp" in fmt: + scale_dtype = MXFP_SCALE_DTYPE + elif "nvfp4" in fmt: + scale_dtype = FP8_DTYPE + else: + scale_dtype = torch.float32 + + scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) + + scale2 = None + if "nvfp4" in fmt or schema.strategy == "tensor_group": + scale2 = ScaleDesc( + dtype=torch.float32, + static=True, + group_shape=GroupShape.PER_TENSOR, + ) + + return QuantKey( + dtype=dtype, + scale=scale, + scale2=scale2, + symmetric=schema.symmetric, + ) + + +# ---- Dispatch for any BaseWeightSchema ------------------------------------ + + +def weight_schema_to_quant_key( + schema: "BaseWeightSchema", +) -> QuantKey: + from vllm.utils.humming import ( + AWQWeightSchema, + BitnetWeightSchema, + CompressedTensorsWeightSchema, + Fp8WeightSchema, + GptOssMxfp4WeightSchema, + GPTQWeightSchema, + HummingWeightSchema, + ModeloptMxfp8WeightSchema, + ModeloptNvfp4WeightSchema, + Mxfp4WeightSchema, + ) + + """Convert any BaseWeightSchema to a QuantKey.""" + if isinstance(schema, HummingWeightSchema): + return _humming_weight_schema_to_quant_key(schema) + + # Schemas with fixed QuantKeys + if isinstance(schema, (Mxfp4WeightSchema, GptOssMxfp4WeightSchema)): + return QuantKey( + dtype=FP4_DTYPE, + scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)), + ) + if isinstance(schema, ModeloptMxfp8WeightSchema): + return QuantKey( + dtype=FP8_DTYPE, + scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)), + ) + if isinstance(schema, ModeloptNvfp4WeightSchema): + return QuantKey( + dtype=FP4_DTYPE, + scale=ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)), + scale2=ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR), + ) + if isinstance(schema, BitnetWeightSchema): + return QuantKey( + dtype=torch.uint8, + scale=ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL), + ) + + # Schemas requiring config inspection + if isinstance(schema, Fp8WeightSchema): + return _fp8_weight_schema_to_quant_key(schema) + if isinstance(schema, AWQWeightSchema): + return _awq_weight_schema_to_quant_key(schema) + if isinstance(schema, GPTQWeightSchema): + return _gptq_weight_schema_to_quant_key(schema) + if isinstance(schema, CompressedTensorsWeightSchema): + return _compressed_tensors_weight_schema_to_quant_key(schema) + + raise TypeError(f"Unsupported weight schema type: {type(schema)}") + + +# ---- HummingInputSchema (post-conversion) ---------------------------------- + + +def _humming_input_schema_to_quant_key( + schema: "HummingInputSchema", +) -> QuantKey | None: + """Convert a HummingInputSchema to a QuantKey. Returns None if + the schema represents unquantized (bf16/fp16) inputs.""" + if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: + return None + + dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype] + + gs = schema.input_scale_group_size + group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN + + scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32 + + scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape) + + return QuantKey(dtype=dtype, scale=scale, symmetric=True) + + +# ---- Checkpoint-format input schemas (pre-conversion) ---------------------- + + +def _resolve_input_quant_key( + origin_a_dtype: "humming_dtypes.DataType", + group_size: int, +) -> QuantKey | None: + from vllm.utils.humming import HummingInputSchema + + """Resolve the actual activation QuantKey after platform fallback.""" + a_dtype = HummingInputSchema().get_fallback_input_dtype(origin_a_dtype) + if a_dtype is None or a_dtype.num_bits >= 16: + return None + + dtype = _HUMMING_TO_QUANT_DTYPE[a_dtype] + gs = group_size if a_dtype == humming_dtypes.float4e2m1 else 0 + group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN + scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32 + + scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape) + return QuantKey(dtype=dtype, scale=scale, symmetric=True) + + +def _compressed_tensors_input_schema_to_quant_key( + schema: "CompressedTensorsInputSchema", +) -> QuantKey | None: + type_bits_to_dtype = { + ("float", 8): humming_dtypes.float8e4m3, + ("float", 4): humming_dtypes.float4e2m1, + ("int", 8): humming_dtypes.int8, + ("int", 4): humming_dtypes.int4, + } + origin = type_bits_to_dtype.get((schema.type, schema.num_bits)) + if origin is None: + return None + return _resolve_input_quant_key(origin, schema.group_size) + + +# ---- Dispatch for any BaseInputSchema ------------------------------------- + + +def input_schema_to_quant_key( + schema: "BaseInputSchema", +) -> QuantKey | None: + from vllm.utils.humming import ( + CompressedTensorsInputSchema, + Fp8InputSchema, + HummingInputSchema, + ModeloptNvfp4InputSchema, + ) + + """Convert any BaseInputSchema to a QuantKey. Returns None if + the schema represents unquantized (bf16/fp16) inputs.""" + if isinstance(schema, HummingInputSchema): + return _humming_input_schema_to_quant_key(schema) + + if isinstance(schema, Fp8InputSchema): + return _resolve_input_quant_key(humming_dtypes.float8e4m3, 0) + + if isinstance(schema, ModeloptNvfp4InputSchema): + return _resolve_input_quant_key( + humming_dtypes.float8e4m3, + schema.group_size, + ) + + if isinstance(schema, CompressedTensorsInputSchema): + return _compressed_tensors_input_schema_to_quant_key(schema) + + raise TypeError(f"Unsupported input schema type: {type(schema)}") def humming_is_layer_skipped(config: dict[str, Any], prefix: str): @@ -24,8 +380,9 @@ def humming_is_layer_skipped(config: dict[str, Any], prefix: str): keys = ["ignored_layers", "ignore", "modules_to_not_convert"] ignored_layers: list[str] = [] for key in keys: - ignored_layers = config.get(key, []) or [] - if not ignored_layers: + candidate = config.get(key, []) or [] + if candidate: + ignored_layers = candidate break if any(module_name in prefix for module_name in ignored_layers): @@ -79,6 +436,12 @@ def convert_linear_layer_to_humming_standard( def prepare_humming_layer(layer: LinearBase, quant_config: dict): + from vllm.utils.humming import ( + BaseWeightSchema, + HummingInputSchema, + HummingMethod, + ) + weight_schema = BaseWeightSchema.from_config(quant_config) input_schema = HummingInputSchema() @@ -140,90 +503,59 @@ def prepare_humming_layer(layer: LinearBase, quant_config: dict): layer.compute_config = json.dumps(compute_config) -def prepare_humming_moe_layer(layer: RoutedExperts, quant_config: dict): - weight_schema = BaseWeightSchema.from_config(quant_config) - input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {} - if humming_is_layer_skipped(input_quant_config, layer.layer_name): - input_schema = HummingInputSchema() +def make_humming_moe_quant_config( + quant_dtype: torch.dtype | str | None, + weight_dtype: torch.dtype | str | None, + weight_group_shape: GroupShape | None = None, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + w1_gscale: torch.Tensor | None = None, + w2_gscale: torch.Tensor | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + if quant_dtype is None: + a_quant_desc = FusedMoEQuantDesc(dtype=None) else: - # TODO: read input_quant_config from quant_config - input_schema = HummingInputSchema.from_config(input_quant_config) + shape = GroupShape(row=1, col=-1) + a_quant_desc = FusedMoEQuantDesc(dtype=quant_dtype, shape=shape) - is_gated = layer.activation.is_gated - shape_config = { - "w13": ( - layer.moe_config.intermediate_size_per_partition * 2, - layer.moe_config.hidden_dim, - ), - "w2": ( - layer.moe_config.hidden_dim, - layer.moe_config.intermediate_size_per_partition * (1 if is_gated else 2), - ), - } + w1_quant_desc = FusedMoEQuantDesc( + dtype=weight_dtype, + shape=weight_group_shape, + scale=w1_scale, + alpha_or_gscale=w1_gscale, + zp=w1_zp, + bias=w1_bias, + ) - layer.weight_schemas = {} - layer.input_schemas = {} + w2_quant_desc = FusedMoEQuantDesc( + dtype=weight_dtype, + shape=weight_group_shape, + scale=w2_scale, + alpha_or_gscale=w2_gscale, + zp=w2_zp, + bias=w2_bias, + ) - for sublayer_name in shape_config: - # Step 1: convert weight to humming standard format - tensors: dict[str, torch.Tensor] = dict( - (key.removeprefix(sublayer_name + "_"), value) - for key, value in layer.state_dict().items() - if key.startswith(sublayer_name + "_") - ) - - shape_n, shape_k = shape_config[sublayer_name] - shape_n_stacks = [shape_n] - shape_k_stacks = [shape_k] - if sublayer_name == "w13": - shape_n_stacks = [shape_n // 2] * 2 - - weight_schema_new, tensors = weight_schema.convert_humming( - tensors=tensors, - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - num_experts=layer.local_num_experts, - param_dtype=layer.params_dtype, - ) - - layer.weight_schemas[sublayer_name] = weight_schema_new - layer.input_schemas[sublayer_name] = input_schema - - for name, _ in list(layer.named_parameters()): - if not name.startswith(sublayer_name + "_"): - continue - delattr(layer, name) - - for name, tensor in tensors.items(): - name = f"{sublayer_name}_{name}" - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - - # Step 2: transform weight (humming standard format) for forwarding - HummingMethod.prepare_layer_meta( - layer=layer, - shape_n=shape_n, - shape_k=shape_k, - pad_n_to_multiple=256, - pad_k_to_multiple=128, - input_schema=input_schema, - weight_schema=weight_schema_new, - has_bias=layer.moe_config.has_bias, - num_experts=layer.num_experts, - torch_dtype=layer.params_dtype, - sublayer_name=sublayer_name, - ) - - HummingMethod.transform_humming_layer(layer, sublayer_name=sublayer_name) - - if not hasattr(layer, "locks"): - device = layer.w13_weight.device - locks = torch.zeros(1024, dtype=torch.int32, device=device) - layer.register_buffer("locks", locks) + return FusedMoEQuantConfig( + _a1=a_quant_desc, + _a2=a_quant_desc, + _w1=w1_quant_desc, + _w2=w2_quant_desc, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) def get_humming_moe_quant_config( - layer: RoutedExperts, + layer: "RoutedExperts", gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, @@ -231,12 +563,10 @@ def get_humming_moe_quant_config( input_schema = layer.input_schemas["w13"] weight_schema = layer.weight_schemas["w13"] - a_dtype = input_schema.a_dtype - if a_dtype is None or a_dtype.num_bits == 16: - a_quant_desc = FusedMoEQuantDesc(dtype=None) + if input_schema.a_dtype is None or input_schema.a_dtype.num_bits == 16: + q_dtype = None else: - shape = GroupShape(row=1, col=-1) - a_quant_desc = FusedMoEQuantDesc(dtype=str(a_dtype), shape=shape) + q_dtype = str(input_schema.a_dtype) weight_scale_group_size = weight_schema.weight_scale_group_size weight_scale_group_size_n = weight_schema.weight_scale_group_size_n @@ -251,30 +581,437 @@ def get_humming_moe_quant_config( else: weight_group_shape = GroupShape(row=weight_scale_group_size, col=1) - w1_quant_desc = FusedMoEQuantDesc( - dtype=str(weight_schema.b_dtype), - shape=weight_group_shape, - scale=getattr(layer, "w13_weight_scale", None), - alpha_or_gscale=getattr(layer, "w13_global_scale", None), - zp=getattr(layer, "w13_zero_point", None), - bias=getattr(layer, "w13_bias", None), + return make_humming_moe_quant_config( + quant_dtype=q_dtype, + weight_dtype=str(weight_schema.b_dtype), + weight_group_shape=weight_group_shape, + w1_scale=getattr(layer, "w13_weight_scale", None), + w1_gscale=getattr(layer, "w13_global_scale", None), + w1_zp=getattr(layer, "w13_zero_point", None), + w1_bias=getattr(layer, "w13_bias", None), + w2_scale=getattr(layer, "w2_weight_scale", None), + w2_gscale=getattr(layer, "w2_global_scale", None), + w2_zp=getattr(layer, "w2_zero_point", None), + w2_bias=getattr(layer, "w2_bias", None), ) - w2_quant_desc = FusedMoEQuantDesc( - dtype=str(weight_schema.b_dtype), - shape=weight_group_shape, - scale=getattr(layer, "w2_weight_scale", None), - alpha_or_gscale=getattr(layer, "w2_global_scale", None), - zp=getattr(layer, "w2_zero_point", None), - bias=getattr(layer, "w2_bias", None), + +def select_humming_moe_experts( + config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, +) -> type[mk.FusedMoEExperts] | None: + """ + Select the primary Humming MoE Experts class + Note: Shape-specific fallbacks may still occur at runtime. + """ + + if not has_humming(): + return None + + # NOTE: the kernels are selected in the following order. + AVAILABLE_EXPERTS: list[type[mk.FusedMoEExperts]] = [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + + # NOTE(rob): We need to peak into the P/F selection to determine + # if we are using the batched or standard expert format, which + # if not ideal. Once we unify TP + DP/EP, we can select P/F first. + activation_format = ( + mk.FusedMoEActivationFormat.BatchedExperts + if config.moe_parallel_config.use_batched_activation_format + else mk.FusedMoEActivationFormat.Standard ) - return FusedMoEQuantConfig( - _a1=a_quant_desc, - _a2=a_quant_desc, - _w1=w1_quant_desc, - _w2=w2_quant_desc, - gemm1_alpha=gemm1_alpha, - gemm1_beta=gemm1_beta, - gemm1_clamp_limit=gemm1_clamp_limit, + def _make_log_backend(experts_cls: type[mk.FusedMoEExperts]): + return f"Using {experts_cls.__name__} Humming MoE backend." + + def _make_log_unsupported( + experts_cls: type[mk.FusedMoEExperts], reason: str | None + ) -> str: + if reason: + return ( + f"Humming MoE experts {experts_cls.__name__} does not support the " + f"deployment configuration since {reason}." + ) + else: + return ( + f"Humming MoE experts '{experts_cls.__name__}' does not support the " + "deployment configuration." + ) + + for k_cls in AVAILABLE_EXPERTS: + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, + ) + if supported: + logger.info_once(_make_log_backend(k_cls)) + return k_cls + else: + logger.debug_once(_make_log_unsupported(k_cls, reason)) + + return None + + +def make_humming_moe_kernel( + moe_quant_config: FusedMoEQuantConfig, + moe_config: FusedMoEConfig, + experts_cls: type[mk.FusedMoEExperts], + layer: "RoutedExperts", + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, +) -> mk.FusedMoEKernel: + # Create Prepare/Finalize. + prepare_finalize = maybe_make_prepare_finalize( + moe=moe_config, + quant_config=moe_quant_config, + routing_tables=routing_tables, + allow_new_interface=True, + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), ) + assert prepare_finalize is not None + + logger.info_once("Using %s", prepare_finalize.__class__.__name__) + + extra_args: dict[str, Any] = {"layer": layer} + + # Create Experts. + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: + max_num_tokens = prepare_finalize.max_num_tokens_per_rank() + assert max_num_tokens is not None + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_args, + ) + else: + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + **extra_args, + ) + + kernel = mk.FusedMoEKernel( + prepare_finalize, + experts, + ) + + return kernel + + +def _extract_sublayer_tensors( + layer: "RoutedExperts", + sublayer_name: str, +) -> dict[str, torch.Tensor]: + """Extract tensors for a specific sublayer from the layer's state dict.""" + return dict( + (key.removeprefix(sublayer_name + "_"), value) + for key, value in layer.state_dict().items() + if key.startswith(sublayer_name + "_") + ) + + +def _replace_layer_parameters( + layer: "RoutedExperts", + sublayer_name: str, + tensors: dict[str, torch.Tensor], + preserve_bias: bool = False, +) -> None: + """ + Replace layer parameters for a sublayer with new tensors. + + Args: + layer: The RoutedExperts layer + sublayer_name: Name of the sublayer (e.g., "w13", "w2") + tensors: Dict of parameter name to tensor + preserve_bias: If True, don't delete bias parameters + """ + # Delete old parameters + for name, _ in list(layer.named_parameters()): + if not name.startswith(sublayer_name + "_"): + continue + if preserve_bias and name == sublayer_name + "_bias": + continue + delattr(layer, name) + + # Set new parameters + for name, tensor in tensors.items(): + param_name = f"{sublayer_name}_{name}" + param = torch.nn.Parameter(tensor, requires_grad=False) + setattr(layer, param_name, param) + + +def _convert_sublayer_to_humming( + layer: "RoutedExperts", + sublayer_name: str, + shape_n: int, + shape_k: int, + weight_schema: Any, + input_schema: Any, + num_experts: int, + param_dtype: torch.dtype, +) -> tuple[Any, Any]: + """ + Convert a sublayer's weights from checkpoint format to Humming format. + + Returns: + Tuple of (converted_weight_schema, converted_input_schema) + """ + from humming.schema import HummingWeightSchema + + if isinstance(weight_schema, HummingWeightSchema): + # Already in Humming format + return weight_schema, input_schema + + tensors = _extract_sublayer_tensors(layer, sublayer_name) + + shape_k_stacks = [shape_k] + shape_n_stacks = [shape_n] + if sublayer_name == "w13": + shape_n_stacks = [shape_n // 2] * 2 + + converted_weight_schema, converted_tensors = weight_schema.convert_humming( + tensors=tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=param_dtype, + num_experts=num_experts, + ) + + converted_input_schema, _ = input_schema.convert_humming( + tensors=converted_tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=param_dtype, + num_experts=num_experts, + ) + + _replace_layer_parameters(layer, sublayer_name, converted_tensors) + + return converted_weight_schema, converted_input_schema + + +def _prepare_and_transform_sublayer( + layer: "RoutedExperts", + sublayer_name: str, + shape_n: int, + shape_k: int, + weight_schema: Any, + input_schema: Any, + has_bias: bool, + num_experts: int, + param_dtype: torch.dtype, +) -> None: + """ + Prepare layer metadata and transform weights for a sublayer. + + This calls Humming's prepare_layer_meta and transform_humming_layer. + """ + from humming.layer import HummingMethod + + HummingMethod.prepare_layer_meta( + layer=layer, + shape_n=shape_n, + shape_k=shape_k, + pad_n_to_multiple=256, + pad_k_to_multiple=128, + input_schema=input_schema, + weight_schema=weight_schema, + has_bias=has_bias, + num_experts=num_experts, + torch_dtype=param_dtype, + sublayer_name=sublayer_name, + ) + + HummingMethod.transform_humming_layer(layer, sublayer_name=sublayer_name) + + +def _process_single_sublayer( + layer: "RoutedExperts", + sublayer_name: str, + shape_n: int, + shape_k: int, + weight_schema: Any, + input_schema: Any, + has_bias: bool, + num_experts: int, + param_dtype: torch.dtype, + force_weight_schema: Any | None = None, +) -> tuple[Any, Any]: + """ + Process a single sublayer: convert, optionally requant, prepare, and transform. + + This combines the common logic from convert_to_humming_moe_kernel_format + for processing a single sublayer. + + Args: + layer: The RoutedExperts layer + sublayer_name: Name of the sublayer (e.g., "w13", "w2") + shape_n: Output dimension size + shape_k: Input dimension size + weight_schema: Initial weight quantization schema + input_schema: Initial input quantization schema + has_bias: Whether the layer has bias terms + num_experts: Number of experts + param_dtype: Parameter data type + force_weight_schema: Optional schema to force requantization to + + Returns: + Tuple of (final_weight_schema, final_input_schema) + """ + from humming.schema import HummingWeightSchema + + # Step 1: Convert from checkpoint format to humming format if needed + current_weight_schema, current_input_schema = _convert_sublayer_to_humming( + layer=layer, + sublayer_name=sublayer_name, + shape_n=shape_n, + shape_k=shape_k, + weight_schema=weight_schema, + input_schema=input_schema, + num_experts=num_experts, + param_dtype=param_dtype, + ) + + # Step 2: Force requant if needed + assert isinstance(current_weight_schema, HummingWeightSchema) + if force_weight_schema is not None and current_weight_schema != force_weight_schema: + tensors = _extract_sublayer_tensors(layer, sublayer_name) + + tensors = current_weight_schema.requant_tensors( + tensors=tensors, + target_weight_schema=force_weight_schema, + param_dtype=param_dtype, + ) + + current_weight_schema = force_weight_schema + _replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True) + del tensors + + # Step 3: Prepare layer metadata and transform weights + _prepare_and_transform_sublayer( + layer=layer, + sublayer_name=sublayer_name, + shape_n=shape_n, + shape_k=shape_k, + weight_schema=current_weight_schema, + input_schema=current_input_schema, + has_bias=has_bias, + num_experts=num_experts, + param_dtype=param_dtype, + ) + + return current_weight_schema, current_input_schema + + +def convert_to_humming_moe_kernel_format( + layer: "RoutedExperts", + quant_config: dict | None = None, + sublayer_configs: dict[str, Any] | None = None, + weight_schema: Any | None = None, + input_schema: Any | None = None, + force_weight_schema: Any | None = None, +) -> None: + """ + Convert MoE weights from checkpoint format to Humming kernel format. + + This function processes weights for each sublayer (w13, w2) by: + 1. Converting from checkpoint format to humming format if needed + 2. Force requanting if a different quantization schema is specified + 3. Preparing layer metadata for the Humming kernel + 4. Transforming weights for inference + + Args: + layer: The RoutedExperts layer containing weights to process + quant_config: Optional quantization config dict. Required if weight_schema + or input_schema are None. Used to build schemas via + BaseWeightSchema.from_config(). + sublayer_configs: Optional configuration dict for each sublayer (w13, w2). + Each config must have "shape_n" and "shape_k" keys. + If None, configs are built from layer.moe_config properties. + weight_schema: Optional initial weight quantization schema. + If None, built from quant_config. + input_schema: Optional initial input quantization schema. + If None, built from quant_config or env vars. + force_weight_schema: Optional schema to force requantization to + + Side effects: + - Modifies layer parameters in place + - Sets layer.weight_schemas and layer.input_schemas + """ + + # Build schemas from quant_config if not provided + has_bias = layer.moe_config.has_bias + num_experts = layer.moe_config.num_local_experts + param_dtype = layer.params_dtype + + if weight_schema is None or input_schema is None: + if quant_config is None: + raise ValueError( + "Must provide either weight_schema/input_schema or quant_config" + ) + + from humming.layer import HummingInputSchema + from humming.schema import BaseWeightSchema + + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + humming_is_layer_skipped, + ) + + if weight_schema is None: + weight_schema = BaseWeightSchema.from_config(quant_config) + + if input_schema is None: + input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {} + if humming_is_layer_skipped(input_quant_config, layer.layer_name): + input_schema = HummingInputSchema() + else: + # TODO: read input_quant_config from quant_config + input_schema = HummingInputSchema.from_config(input_quant_config) + + # Build sublayer configs from layer properties if not provided + if sublayer_configs is None: + is_gated = layer.moe_config.activation.is_gated + sublayer_configs = { + "w13": { + "shape_n": layer.moe_config.intermediate_size_per_partition * 2, + "shape_k": layer.moe_config.hidden_dim, + }, + "w2": { + "shape_n": layer.moe_config.hidden_dim, + "shape_k": layer.moe_config.intermediate_size_per_partition + * (1 if is_gated else 2), + }, + } + + layer.weight_schemas = {} + layer.input_schemas = {} + + for sublayer_name, configs in sublayer_configs.items(): + final_weight_schema, final_input_schema = _process_single_sublayer( + layer=layer, + sublayer_name=sublayer_name, + shape_n=configs["shape_n"], + shape_k=configs["shape_k"], + weight_schema=weight_schema, + input_schema=input_schema, + has_bias=has_bias, + num_experts=num_experts, + param_dtype=param_dtype, + force_weight_schema=force_weight_schema, + ) + + layer.weight_schemas[sublayer_name] = final_weight_schema + layer.input_schemas[sublayer_name] = final_input_schema + + if not hasattr(layer, "locks"): + device = layer.w13_weight.device + locks = torch.zeros(1024, dtype=torch.int32, device=device) + layer.register_buffer("locks", locks) diff --git a/vllm/utils/humming.py b/vllm/utils/humming.py index b8d9445c3f3..bdd519bd8c5 100644 --- a/vllm/utils/humming.py +++ b/vllm/utils/humming.py @@ -15,6 +15,7 @@ _EXPORTS: dict[str, str] = { "dtypes": "humming.dtypes", "DataType": "humming.dtypes:DataType", "GemmType": "humming.config:GemmType", + "WeightScaleType": "humming.config:WeightScaleType", "HummingMethod": "humming.layer:HummingMethod", "HummingLayerMeta": "humming.layer:HummingLayerMeta", "BaseInputSchema": "humming.schema:BaseInputSchema", @@ -22,6 +23,18 @@ _EXPORTS: dict[str, str] = { "HummingInputSchema": "humming.schema:HummingInputSchema", "HummingWeightSchema": "humming.schema:HummingWeightSchema", "quantize_weight": "humming.utils.weight:quantize_weight", + "AWQWeightSchema": "humming.schema:AWQWeightSchema", + "BitnetWeightSchema": "humming.schema:BitnetWeightSchema", + "ModeloptMxfp8WeightSchema": "humming.schema.modelopt:ModeloptMxfp8WeightSchema", + "ModeloptNvfp4InputSchema": "humming.schema.modelopt:ModeloptNvfp4InputSchema", + "ModeloptNvfp4WeightSchema": "humming.schema.modelopt:ModeloptNvfp4WeightSchema", + "CompressedTensorsInputSchema": "humming.schema:CompressedTensorsInputSchema", + "CompressedTensorsWeightSchema": "humming.schema:CompressedTensorsWeightSchema", + "Fp8InputSchema": "humming.schema:Fp8InputSchema", + "Fp8WeightSchema": "humming.schema.fp8:Fp8WeightSchema", + "Mxfp4WeightSchema": "humming.schema:Mxfp4WeightSchema", + "GptOssMxfp4WeightSchema": "humming.schema:GptOssMxfp4WeightSchema", + "GPTQWeightSchema": "humming.schema:GPTQWeightSchema", } From 6185d73882c0cdfd9ee13cea16a9b50d2b5267be Mon Sep 17 00:00:00 2001 From: Blas Rodriguez Irizar Date: Mon, 29 Jun 2026 14:46:33 +0100 Subject: [PATCH 116/138] [Rust Frontend] Keep literal "null" string for string-typed tool params (#46827) Signed-off-by: Blas Rodriguez Irizar --- .../src/tool/deepseek_dsml/deepseek_v32.rs | 2 +- rust/src/parser/src/tool/parameters.rs | 63 ++++++++++++++----- 2 files changed, 48 insertions(+), 17 deletions(-) diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs index 7d3432f9e5c..e4f5c58ee0e 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs @@ -154,7 +154,7 @@ mod tests { "flag": true, "payload": { "nested": true }, "items": [1, 2], - "empty": null, + "empty": "null", }) ); } diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index f5661456e3e..f9abb50f8b1 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -166,7 +166,13 @@ impl JsonParamType { // Typically, these types are already handled by checking the "type" field, but // we can also infer them from their characteristic fields if "type" is missing. - if schema.contains_key("enum") { + if let Some(values) = schema.get("enum").and_then(Value::as_array) { + // Enum values are treated as strings, except that a `null` member + // makes the parameter nullable (mirrors Python's enum type + // inference), so a literal "null" coerces to JSON null. + if values.iter().any(Value::is_null) { + return Some(Self::one_of(vec![Self::String, Self::Null])); + } return Some(Self::String); } if schema.contains_key("items") { @@ -277,9 +283,12 @@ impl JsonParamType { /// Convert one parameter input to a normalized JSON value. fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value { - // For literal `null`, always convert to JSON null value. + // Coerce the literal text `null` to JSON null, except for `string`-typed + // params, where it must stay the string "null": a model emitting the literal + // text "null" for a string field means the string, not a missing value. if let ParamInput::Text(value) = input && value.eq_ignore_ascii_case("null") + && param_type != Some(&JsonParamType::String) { return Value::Null; } @@ -685,21 +694,43 @@ mod tests { } #[test] - fn convert_params_preserves_null_for_known_param() { - let schemas = ToolSchemas::from_tools(&[test_tool( - "convert", - json!({ - "type": "object", - "properties": { - "value": { "type": "string" } - } - }), - )]); + fn string_param_preserves_literal_null_text() { + // A `string`-typed param whose value is the literal text "null"/"NULL" + // must stay a string (the original case is preserved), rather than being + // coerced to JSON null. Non-string types keep coercing "null" to null. + let params = ToolSchema::from_schema(&json!({ + "type": "object", + "properties": { + "name": { "type": "string" }, + "count": { "type": "integer" }, + "anything": {} + } + })); - let converted = schemas - .convert_params_with_schema("convert", vec![("value".to_string(), "NULL".to_string())]); + assert_eq!(params.convert("name", text("null")), json!("null")); + assert_eq!(params.convert("name", text("NULL")), json!("NULL")); + // Non-string and schema-less params are unchanged: "null" -> null. + assert_eq!(params.convert("count", text("null")), json!(null)); + assert_eq!(params.convert("anything", text("null")), json!(null)); + } - assert_eq!(converted.get("value"), Some(&json!(null))); + #[test] + fn nullable_enum_param_coerces_literal_null() { + // An enum that includes `null` admits a null value, so a literal "null" + // must coerce to JSON null (matching Python's `extract_types_from_schema`, + // which infers `null` from the enum values), while a non-null enum keeps + // "null" as a string. + let params = ToolSchema::from_schema(&json!({ + "type": "object", + "properties": { + "mode": { "enum": [null, "auto"] }, + "color": { "enum": ["red", "green"] } + } + })); + + assert_eq!(params.convert("mode", text("null")), json!(null)); + assert_eq!(params.convert("mode", text("auto")), json!("auto")); + assert_eq!(params.convert("color", text("null")), json!("null")); } #[test] @@ -841,7 +872,7 @@ mod tests { "user_id": 42, "urgent": true, "note": "Please leave at front desk.", - "nil": null, + "nil": "NULL", "shipping": { "city": "Singapore", "zip": 18956 From 0ca39c4f1fc450339f57ceca6bddc2af1abe84a5 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Mon, 29 Jun 2026 10:00:31 -0400 Subject: [PATCH 117/138] [Bugfix] Capture final-layer aux hidden state in deepseek_v2 backbone (#46973) Signed-off-by: mgoin Co-authored-by: Claude Opus 4.8 (1M context) --- vllm/model_executor/models/deepseek_v2.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 2c6b075ae74..144ff3971a2 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1470,6 +1470,9 @@ class DeepseekV2Model(nn.Module): [self.hidden_size, self.hidden_size], dim=-1 ) + if self.end_layer in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states + residual) + hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states From 49e28e8e91ad9ed102e88e50c190686408be5552 Mon Sep 17 00:00:00 2001 From: "Xiaohong (Sean) Chen" Date: Mon, 29 Jun 2026 10:54:15 -0400 Subject: [PATCH 118/138] [Kernel][Helion][1/N] Add Helion kernel for fused_qk_norm_rope (#44010) Signed-off-by: Sean Chen --- .../kernels/helion/test_fused_qk_norm_rope.py | 261 ++ .../fused_qk_norm_rope/nvidia_b200.json | 2612 ++++++++++++++++ .../fused_qk_norm_rope/nvidia_h100.json | 2722 +++++++++++++++++ vllm/kernels/helion/ops/fused_qk_norm_rope.py | 316 ++ 4 files changed, 5911 insertions(+) create mode 100644 tests/kernels/helion/test_fused_qk_norm_rope.py create mode 100644 vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_b200.json create mode 100644 vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_h100.json create mode 100644 vllm/kernels/helion/ops/fused_qk_norm_rope.py diff --git a/tests/kernels/helion/test_fused_qk_norm_rope.py b/tests/kernels/helion/test_fused_qk_norm_rope.py new file mode 100644 index 00000000000..19d2fc9b5a6 --- /dev/null +++ b/tests/kernels/helion/test_fused_qk_norm_rope.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the fused_qk_norm_rope helion kernel + +Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`. +""" + +from typing import Any + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tests.kernels.helion.utils import skip_if_platform_unsupported +from vllm.benchmarks.lib.utils import default_vllm_config +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.config_manager import ConfigManager +from vllm.kernels.helion.ops.fused_qk_norm_rope import ( + _pick_cache, + baseline, + fused_qk_norm_rope, + pick_config, +) +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.utils.import_utils import has_helion + +if not has_helion(): + pytest.skip( + "Helion is not installed. Install with: pip install vllm[helion]", + allow_module_level=True, + ) + + +@default_vllm_config() +def _generate_fake_input( + num_tokens: int, num_q_heads: int, num_kv_heads: int +) -> tuple[Any, ...]: + with FakeTensorMode(): + head_dim = 128 + eps = 1e-6 + is_neox = True + rotary_ratio = 1.0 + device = "cuda" + dtype = torch.bfloat16 + total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim + qkv = torch.randn(num_tokens, total_dim, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.long, device=device) + q_weight = torch.normal( + mean=1.0, + std=1.0, + size=(head_dim,), + dtype=qkv.dtype, + device=device, + ) + k_weight = torch.normal( + mean=1.0, + std=1.0, + size=(head_dim,), + dtype=qkv.dtype, + device=device, + ) + rotary_dim = int(head_dim * rotary_ratio) + rope = RotaryEmbedding( + head_size=head_dim, + rotary_dim=rotary_dim, + max_position_embeddings=4096, + base=10000.0, + is_neox_style=is_neox, + dtype=dtype, + ).to(device) + args = ( + qkv, + num_q_heads, + num_kv_heads, + num_kv_heads, + head_dim, + eps, + q_weight, + k_weight, + rope.cos_sin_cache, + is_neox, + positions.view(-1), + ) + return args + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestFusedQkNormRopeConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}), + CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}), + ] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"q_heads": 4096, "kv_heads": 128, "num_tokens": 16} + ) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}), + CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}), + CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}), + CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}), + CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}), + CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}), + CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}), + CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(20, 3000, 70) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"q_heads": 2048, "kv_heads": 64, "num_tokens": 32} + ) + + def test_config_picker_no_configs(self): + config_keys: list[dict] = [] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key is None + + def test_config_picker_fallback_to_largest(self): + config_keys = [ + CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}), + CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}), + CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}), + CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}), + CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}), + CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}), + CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}), + CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(64, 8192, 256) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"q_heads": 4096, "kv_heads": 128, "num_tokens": 32} + ) + + +class TestFusedQkNormRopeCorrectness: + @pytest.mark.parametrize( + "num_heads, num_kv_heads, head_dim", [(16, 4, 128), (64, 8, 128)] + ) + @pytest.mark.parametrize("num_tokens", [1, 7, 1024, 1025]) + @pytest.mark.parametrize("is_neox", [False, True]) + @pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25]) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + @default_vllm_config() + def test_fused_qk_norm_rope( + self, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_tokens: int, + is_neox: bool, + rotary_ratio: float, + dtype: torch.dtype, + ): + skip_if_platform_unsupported("fused_qk_norm_rope") + + torch.manual_seed(42) + eps = 1e-6 + device = "cuda" + total_dim = (num_heads + 2 * num_kv_heads) * head_dim + ref_qkv = torch.empty( + num_tokens, total_dim, dtype=dtype, device=device + ).uniform_(-0.1, 0.1) + ops_qkv = ref_qkv.clone() + positions = torch.arange(num_tokens, dtype=torch.long, device=device) + q_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2) + k_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2) + rotary_dim = int(head_dim * rotary_ratio) + rope = RotaryEmbedding( + head_size=head_dim, + rotary_dim=rotary_dim, + max_position_embeddings=40960, + base=10000.0, + is_neox_style=is_neox, + dtype=dtype, + ).to(device) + + baseline( + ref_qkv, + num_heads, + num_kv_heads, + num_kv_heads, + head_dim, + eps, + q_weight, + k_weight, + rope.cos_sin_cache, + is_neox, + positions.view(-1), + ) + + fused_qk_norm_rope( + ops_qkv, + num_heads, + num_kv_heads, + num_kv_heads, + head_dim, + eps, + q_weight, + k_weight, + rope.cos_sin_cache, + is_neox, + positions.view(-1), + ) + + if dtype == torch.bfloat16: + atol = 5e-2 + rtol = 5e-2 + else: + atol = 1e-2 + rtol = 1e-2 + + torch.testing.assert_close( + ref_qkv, + ops_qkv, + atol=atol, + rtol=rtol, + ) + + +class TestFusedQkNormRopeIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "fused_qk_norm_rope" in registered_kernels + + kernel_wrapper = registered_kernels["fused_qk_norm_rope"] + assert kernel_wrapper.op_name == "fused_qk_norm_rope" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["qkv"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("fused_qk_norm_rope") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["fused_qk_norm_rope"] + fake_impl = kernel_wrapper._fake_impl + + args = _generate_fake_input(16, 4096, 128) + assert fake_impl(*args) is None diff --git a/vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_b200.json b/vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_b200.json new file mode 100644 index 00000000000..cf806f0a4d1 --- /dev/null +++ b/vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_b200.json @@ -0,0 +1,2612 @@ +[ + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "", + "last", + "first", + "", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "last", + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "first", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "", + "last", + "first", + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "first", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "first", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "first", + "last", + "", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last", + "first", + "first", + "first", + "", + "first" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "first", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "", + "", + "last", + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "first", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first", + "first", + "last", + "last", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "", + "", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "first", + "", + "first", + "first", + "", + "first" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "", + "", + "last", + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "", + "", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last", + "", + "last", + "last", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "", + "", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "", + "", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "first", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "first", + "", + "last", + "first", + "first", + "first" + ], + "num_warps": 2, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "", + "", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "", + "last", + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "", + "last", + "first", + "", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "", + "", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "first", + "", + "first", + "first", + "", + "first" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last", + "", + "last", + "last", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "last", + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 4 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "last", + "", + "first", + "first", + "", + "first", + "last", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "last", + "", + "", + "last", + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 32 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "last", + "", + "", + "last", + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 32 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "first", + "", + "first", + "" + ], + "num_warps": 1, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 32 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "last" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + false + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_h100.json b/vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_h100.json new file mode 100644 index 00000000000..825c63ca357 --- /dev/null +++ b/vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_h100.json @@ -0,0 +1,2722 @@ +[ + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 4 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "last", + "", + "first" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 2 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "last", + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 0, + 1 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last", + "first", + "first", + "last", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "first", + "last", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "last", + "first", + "first", + "first", + "last", + "" + ], + "num_warps": 1, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "", + "", + "last", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "", + "", + "last", + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last", + "", + "first", + "", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last", + "first", + "last", + "last", + "last", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first", + "last", + "first", + "first", + "last", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat", + "atomic_indexing": [], + "range_warp_specializes": [], + "range_num_stages": [] + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last", + "first", + "last", + "last", + "last", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last", + "first", + "last", + "last", + "last", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "", + "", + "last", + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last", + "first", + "last", + "last", + "last", + "" + ], + "num_warps": 1, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "", + "first", + "last", + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "last", + "", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "", + "first", + "last", + "", + "last", + "last" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat", + "atomic_indexing": [], + "range_warp_specializes": [], + "range_num_stages": [] + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last", + "last", + "first", + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "last", + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "last", + "", + "last", + "last", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last", + "last", + "last", + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "last", + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "last", + "last", + "last", + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "", + "", + "first", + "", + "first", + "first" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "last", + "last", + "last", + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 4 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "first", + "", + "first", + "first", + "last", + "last", + "", + "" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 64, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "first", + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "last", + "", + "last", + "", + "first" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 128, + "maxnreg": 64, + "atomic_indexing": [], + "range_warp_specializes": [] + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "last", + "", + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 64, + "maxnreg": 256 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "", + "first", + "first", + "last", + "", + "", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 64, + "maxnreg": 128 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "", + "first", + "first", + "last", + "", + "", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 64, + "maxnreg": 128 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "last", + "", + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 64, + "maxnreg": 256 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 16, + "kv_heads": 8, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "", + "last", + "last", + "", + "last", + "last", + "", + "" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64, + "atomic_indexing": [], + "range_warp_specializes": [] + } + }, + { + "key": { + "q_heads": 32, + "kv_heads": 8, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + }, + { + "key": { + "q_heads": 64, + "kv_heads": 8, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 128, + "maxnreg": 64 + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/fused_qk_norm_rope.py b/vllm/kernels/helion/ops/fused_qk_norm_rope.py new file mode 100644 index 00000000000..c97ad1e5145 --- /dev/null +++ b/vllm/kernels/helion/ops/fused_qk_norm_rope.py @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import product +from typing import Any + +import torch + +from vllm.kernels.helion.case_key import CaseKey +from vllm.logger import init_logger +from vllm.utils.import_utils import has_helion + +if not has_helion(): + raise ImportError( + "Helion kernel requires helion to be installed. " + "Install it with: pip install helion" + ) + +import helion +import helion.language as hl + +from vllm import ir +from vllm.kernels.helion.register import register_kernel +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding + +logger = init_logger(__name__) + + +def _compute_cos_sin_cache( + max_position_embeddings, rotary_dim, device="cuda", dtype=torch.float +): + inv_freq = 1.0 / ( + 10000 + ** (torch.arange(0, rotary_dim, 2, device=device, dtype=dtype) / rotary_dim) + ) + + t = torch.arange(max_position_embeddings, device=device, dtype=dtype) + + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + cache = torch.cat((cos, sin), dim=-1) + return cache + + +def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: + # TODO(xiaohongchen1991): it is difficult for kernel author to cover + # all input property combination. Currently, dtypes are fixed. We need + # optimization to bucket/skip some combinations + num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + num_heads_pair = [ + (16, 8), + (32, 8), + (64, 8), + ] + head_dim = 128 + in_dtype: torch.dtype = torch.bfloat16 + rotary_ratio = 1.0 + is_neox = True + eps = 1e-6 + device = "cuda" + inputs = {} + + for num_tokens, (num_q_heads, num_kv_heads) in product( + num_tokens_list, num_heads_pair + ): + total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim + qkv = torch.empty( + num_tokens, total_dim, dtype=in_dtype, device=device + ).uniform_(-0.1, 0.1) + positions = torch.arange(num_tokens, dtype=torch.long, device=device) + q_weight = torch.empty(head_dim, dtype=in_dtype, device=device).uniform_( + 0.8, 1.2 + ) + k_weight = torch.empty(head_dim, dtype=in_dtype, device=device).uniform_( + 0.8, 1.2 + ) + rotary_dim = int(head_dim * rotary_ratio) + cos_sin_cache = _compute_cos_sin_cache(40960, rotary_dim) + cos_sin_cache = cos_sin_cache.to(in_dtype) + + config_key = CaseKey( + { + "q_heads": num_q_heads, + "kv_heads": num_kv_heads, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = ( + qkv, + num_q_heads, + num_kv_heads, + num_kv_heads, + head_dim, + eps, + q_weight, + k_weight, + cos_sin_cache, + is_neox, + positions.view(-1), + ) + + return inputs + + +_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} + + +def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: + """Pick the best pre-tuned config for the given input shape. + + Selection strategy: + 1. Find the closest q_heads among available configs + (exact match preferred). + 2. Find the closest kv_heads among available configs + (exact match preferred). + 3. Among the num_tokens values tuned for that q_heads and q_heads, pick + the smallest num_tokens >= the input's num_tokens. If the input is + larger than all available num_tokens, fall back to the largest. + """ + + if not config_keys: + return None + + qkv, q_heads, kv_heads, *_ = args + num_tokens = qkv.shape[0] + + cache_key = (num_tokens, q_heads, kv_heads) + cached = _pick_cache.get(cache_key) + if cached is not None: + return cached + + configs: dict[int, dict[int, list[int]]] = {} + for key in config_keys: + if key.is_default(): + continue + configs.setdefault(key["q_heads"], {}).setdefault(key["kv_heads"], []).append( + key["num_tokens"] + ) + + if not configs: + return None + + best_q_heads = min(configs, key=lambda s: abs(s - q_heads)) + best_kv_heads = min(configs[best_q_heads], key=lambda s: abs(s - kv_heads)) + available_num_tokens = sorted(configs[best_q_heads][best_kv_heads]) + best_num_tokens = next( + (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] + ) + + result = CaseKey( + { + "q_heads": best_q_heads, + "kv_heads": best_kv_heads, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, # [max_position, rotary_dim] + is_neox: bool, + position_ids: torch.Tensor, # [num_tokens], + forced_token_heads_per_warp: int = -1, # dummy +) -> None: + return + + +def baseline( + qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, # [max_position, rotary_dim] + is_neox: bool, + position_ids: torch.Tensor, # [num_tokens], + forced_token_heads_per_warp: int = -1, # dummy +) -> None: + q_size = num_heads_q * head_dim + kv_size = num_heads_k * head_dim + + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim) + q_by_head = ir.ops.rms_norm(q_by_head, q_weight, eps) + q = q_by_head.view(q.shape) + + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim) + k_by_head = ir.ops.rms_norm(k_by_head, k_weight, eps) + k = k_by_head.view(k.shape) + + q, k = RotaryEmbedding.forward_static( + position_ids, q, k, head_dim, cos_sin_cache.shape[1], cos_sin_cache, is_neox + ) + qkv[:, :q_size].copy_(q) + qkv[:, q_size : q_size + kv_size].copy_(k) + + +# Overwrite autotune_baseline_atol and autotune_baseline_rtol +# if too many configs failed due to baseline check during autotuning +@register_kernel( + mutates_args=["qkv"], + config_picker=pick_config, + input_generator=generate_inputs, + fake_impl=fake_impl, + helion_settings=helion.Settings( + autotune_baseline_fn=baseline, + autotune_baseline_atol=5e-2, + autotune_baseline_rtol=5e-2, + ignore_warnings=[helion.exc.TensorOperationInWrapper], + ), +) # type: ignore[misc] +def fused_qk_norm_rope( + qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, # [max_position, rotary_dim] + is_neox: bool, + position_ids: torch.Tensor, # [num_tokens], + forced_token_heads_per_warp: int = -1, # dummy +) -> None: + assert qkv.ndim == 2 + num_tokens = qkv.shape[0] + total_heads = num_heads_q + num_heads_k + num_heads_v + assert qkv.shape[1] == total_heads * head_dim + hl.specialize(qkv.shape[1]) + + assert cos_sin_cache.ndim == 2 + max_position, rotary_dim = cos_sin_cache.shape + hl.specialize(max_position) + hl.specialize(rotary_dim) + assert rotary_dim % 2 == 0 + assert rotary_dim <= head_dim + embed_dim = rotary_dim // 2 + + hl.specialize(num_heads_q) + hl.specialize(num_heads_k) + hl.specialize(num_heads_v) + hl.specialize(head_dim) + + assert position_ids.ndim == 1 and position_ids.shape[0] == num_tokens + hl.specialize(position_ids.shape[0]) + + assert q_weight.ndim == 1 and q_weight.shape[0] == head_dim + hl.specialize(q_weight.shape[0]) + assert k_weight.ndim == 1 and k_weight.shape[0] == head_dim + hl.specialize(k_weight.shape[0]) + + assert qkv.dtype == q_weight.dtype and q_weight.dtype == k_weight.dtype + assert position_ids.dtype == torch.int64 + + assert qkv.is_contiguous() + assert position_ids.is_contiguous() + assert q_weight.is_contiguous() + assert k_weight.is_contiguous() + assert cos_sin_cache.is_contiguous() + + qk_heads = num_heads_q + num_heads_k + + qkv = qkv.view(num_tokens, -1, head_dim) + + for tile_m, tile_gn, tile_n in hl.tile( + [num_tokens, qk_heads, head_dim], block_size=[1, None, head_dim] + ): + x_blk = qkv[tile_m, tile_gn, tile_n].to(dtype=torch.float32) + + rms = x_blk.pow(2).sum(dim=-1) + rms = torch.rsqrt(rms * (1.0 / head_dim) + eps) + + use_q_weight = (tile_gn.index < num_heads_q)[None, :, None] + w_blk = torch.where( + use_q_weight, q_weight[None, None, tile_n], k_weight[None, None, tile_n] + ) + + x_blk = (x_blk * rms[:, :, None]).to(qkv.dtype) * w_blk + + qkv[tile_m, tile_gn, tile_n] = x_blk + + pos_id = position_ids[tile_m] + cos_blk = cos_sin_cache[pos_id, hl.arange(embed_dim)] + sin_blk = cos_sin_cache[pos_id, hl.arange(embed_dim) + embed_dim] + + if is_neox: + x1_offset = hl.arange(embed_dim) + x2_offset = x1_offset + embed_dim + else: + x1_offset = hl.arange(embed_dim) * 2 + x2_offset = x1_offset + 1 + + x1_blk = qkv[tile_m, tile_gn, x1_offset] + x2_blk = qkv[tile_m, tile_gn, x2_offset] + + o1_blk = x1_blk * cos_blk[:, None, :] - x2_blk * sin_blk[:, None, :] + o2_blk = x2_blk * cos_blk[:, None, :] + x1_blk * sin_blk[:, None, :] + + qkv[tile_m, tile_gn, x1_offset] = o1_blk + qkv[tile_m, tile_gn, x2_offset] = o2_blk From 6149187a4cca41e4e16c6461de39a6c33005c361 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 29 Jun 2026 16:54:29 +0200 Subject: [PATCH 119/138] [Kernel] Triton MLA logits workspace (#46819) Signed-off-by: NickLucche --- vllm/v1/attention/backends/mla/triton_mla.py | 91 ++++++++++++++------ 1 file changed, 63 insertions(+), 28 deletions(-) diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index db11cd2845e..3b91c22516d 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -25,13 +25,58 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) logger = init_logger(__name__) +# num_kv_splits selection (shared by forward_mqa and the workspace reservation +# so the two cannot drift). Both are hardware dependent. +_MIN_WORK_PER_SPLIT = 512 +_SPLIT_OCCUPANCY_MULTIPLIER = 2 + + +def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int: + # Power of 2 to avoid excessive kernel instantiations, capped by an SM-based + # maximum (occupancy multiplier allows multiple blocks per SM + # for latency hiding). + ideal_splits = triton.next_power_of_2(max(1, max_seq_len // _MIN_WORK_PER_SPLIT)) + max_splits = sm_count * _SPLIT_OCCUPANCY_MULTIPLIER + return min(ideal_splits, max_splits) + class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): _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._reserve_attn_logits_workspace() + + def _reserve_attn_logits_workspace(self) -> None: + """Pre-size the shared workspace for the decode split-KV attn logits. + + Reserving at the worst case (max_model_len -> max num_kv_splits, + max_num_seqs decode tokens) before warmup/cudagraph capture means the + per-call ``get_simultaneous`` in ``forward_mqa`` never has to grow the + buffer at runtime (which would raise once the workspace is locked). + """ + if not is_workspace_manager_initialized(): + return + # Decode reorder threshold is 1, so decode tokens <= max_num_seqs. + B = self.vllm_config.scheduler_config.max_num_seqs + # DCP all-gathers the query heads before forward_mqa. + q_num_heads = self.num_heads * self.dcp_world_size + max_splits = _compute_num_kv_splits( + self.model_config.max_model_len, + current_platform.num_compute_units(), + ) + lse_dim = self.mla_dims.kv_lora_rank + 1 + current_workspace_manager().get_simultaneous( + ((B, q_num_heads, max_splits, lse_dim), torch.float32), + ) + class TritonMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -166,35 +211,25 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): if envs.VLLM_BATCH_INVARIANT: num_kv_splits = 1 else: - # Minimum work per split - # hardware dependent - min_work_per_split = 512 + num_kv_splits = _compute_num_kv_splits( + attn_metadata.max_seq_len, self._sm_count + ) - ideal_splits = max(1, attn_metadata.max_seq_len // min_work_per_split) - - # use power of 2 to avoid excessive kernel instantiations - ideal_splits = triton.next_power_of_2(ideal_splits) - - # Calculate SM-based maximum splits with occupancy multiplier - # 2-4x allows multiple blocks per SM for latency hiding - # hardware dependent - occupancy_multiplier = 2 - max_splits = self._sm_count * occupancy_multiplier - num_kv_splits = min(ideal_splits, max_splits) - - # TODO(lucas) Allocate ahead of time - attn_logits = torch.empty( - ( - B, - q_num_heads, - num_kv_splits, - # NOTE: the +1 stores the LogSumExp (LSE) that the stage2 - # kernel uses to merge partial attention outputs across splits. - self.kv_lora_rank + 1, - ), - dtype=torch.float32, - device=q.device, - ) + # NOTE: the +1 stores the LogSumExp (LSE) that the stage2 kernel uses to + # merge partial attention outputs across splits. The scratch is served + # from the shared workspace (reserved at max in the metadata builder), so + # there is no per-call allocation on the decode hot path. Fall back to a + # direct allocation when the workspace manager is not initialized (e.g. + # unit tests without a GPUModelRunner). + logits_shape = (B, q_num_heads, num_kv_splits, self.kv_lora_rank + 1) + if is_workspace_manager_initialized(): + (attn_logits,) = current_workspace_manager().get_simultaneous( + (logits_shape, torch.float32), + ) + else: + attn_logits = torch.empty( + logits_shape, dtype=torch.float32, device=q.device + ) # Add a head dim of 1 kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2) From 36bbecd6436d0dd4c7a27fbb09a787e00534d647 Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Mon, 29 Jun 2026 10:54:34 -0400 Subject: [PATCH 120/138] [BugFix] Revert "[KV Offload] Use background thread for mmap / cpu_tensors pinning" (#46958) Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath --- vllm/v1/kv_offload/cpu/gpu_worker.py | 127 ++++++++------------------- 1 file changed, 38 insertions(+), 89 deletions(-) diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 843e1538f90..c8b9915a1e5 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools -import threading import time from collections import deque from dataclasses import dataclass @@ -121,6 +120,36 @@ def compute_sub_block_ptrs( output[:] = flat[skip_count : skip_count + num_sub_blocks] +def pin_mmap_region(region: SharedOffloadRegion) -> None: + """Register the entire mmap as CUDA pinned memory via cudaHostRegister.""" + if not current_platform.is_cuda_alike(): + logger.info( + "Skipping mmap host registration on %s; cudaHostRegister is only " + "available on CUDA/ROCm.", + current_platform.device_name, + ) + return + + rank = region.rank + + base_ptr = region._base.data_ptr() + result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0) + if result.value != 0: + logger.warning( + "cudaHostRegister failed for rank=%d (code=%d) โ€” " + "transfers will still work but may be slower (unpinned DMA)", + rank, + result, + ) + else: + logger.debug( + "cudaHostRegister rank=%d %.2f GB", + rank, + region.total_size_bytes / 1e9, + ) + region.is_pinned = True + + def _new_descriptor_buffers( num_copy_ops: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -150,8 +179,6 @@ class SingleDirectionOffloadingHandler: kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, mmap_region: SharedOffloadRegion | None = None, - pin_thread: threading.Thread | None = None, - manually_pinned_tensors: list[torch.Tensor] | None = None, ): """ Initialize a SingleDirectionOffloadingHandler. @@ -199,8 +226,6 @@ class SingleDirectionOffloadingHandler: # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it) self._mmap_region = mmap_region - self._pin_thread = pin_thread - self._manually_pinned_tensors = manually_pinned_tensors # job_id -> event self._transfer_events: dict[int, torch.Event] = {} # queue of transfers (job_id, stream, event) @@ -433,23 +458,8 @@ class SingleDirectionOffloadingHandler: self._stream_pool.clear() self._event_pool.clear() self._buffer_pool.clear() - - if self._pin_thread is not None: - self._pin_thread.join() - self._pin_thread = None - - if self._manually_pinned_tensors is not None: - for tensor in self._manually_pinned_tensors: - result = torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr()) - if result.value != 0: - logger.warning( - "cudaHostUnregister failed for CPU tensor (code=%d)", - result.value, - ) - self.src_tensors.clear() self.dst_tensors.clear() - if self._mmap_region is not None: self._mmap_region.cleanup() self._mmap_region = None @@ -471,14 +481,12 @@ class CPUOffloadingWorker(OffloadingWorker): mmap_region: SharedOffloadRegion | None = None, ): pin_memory = PIN_MEMORY - self.pin_thread: threading.Thread | None = None - self._manually_pinned_tensors: list[torch.Tensor] = [] - logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors)) - self._mmap_region = mmap_region + if mmap_region is not None and pin_memory: + pin_mmap_region(mmap_region) gpu_tensors: list[torch.Tensor] = [] - self.cpu_tensors: list[torch.Tensor] = [] + cpu_tensors: list[torch.Tensor] = [] for kv_cache_tensor in kv_caches.tensors: gpu_page_size_bytes = kv_cache_tensor.page_size_bytes gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view( @@ -494,13 +502,10 @@ class CPUOffloadingWorker(OffloadingWorker): (num_cpu_blocks, cpu_page_size_bytes), dtype=torch.int8, device="cpu", - # CUDA/ROCm memory is registered asynchronously below. - # Pinning here would block worker initialization; other - # hardware need PyTorch allocation-time pinning. - pin_memory=PIN_MEMORY and not current_platform.is_cuda_alike(), + pin_memory=pin_memory, ) logger.debug( - "torch.zeros tensor %dร—%d (%.2f GB): %.3f s", + "torch.zeros pinned tensor %dร—%d (%.2f GB): %.3f s", num_cpu_blocks, cpu_page_size_bytes, num_cpu_blocks * cpu_page_size_bytes / 1e9, @@ -508,81 +513,25 @@ class CPUOffloadingWorker(OffloadingWorker): ) gpu_tensors.append(gpu_tensor) - self.cpu_tensors.append(cpu_tensor) - - if pin_memory: - if not current_platform.is_cuda_alike(): - logger.info( - "Skipping host registration on %s; cudaHostRegister is only " - "available on CUDA/ROCm.", - current_platform.device_name, - ) - else: - self.pin_thread = threading.Thread( - target=self._pin_cpu_tensors, - name="CPUTensorPinThread", - ) - self.pin_thread.start() - logger.info("Starting to pin memory in background...") + cpu_tensors.append(cpu_tensor) self._store_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, - cpu_tensors=self.cpu_tensors, + cpu_tensors=cpu_tensors, block_size_factor=block_size_factor, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=True, mmap_region=mmap_region, - pin_thread=self.pin_thread, - manually_pinned_tensors=self._manually_pinned_tensors, ) self._load_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, - cpu_tensors=self.cpu_tensors, + cpu_tensors=cpu_tensors, block_size_factor=block_size_factor, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=False, ) - def _pin_cpu_tensors(self) -> None: - """Register the CPU offload memory as CUDA pinned memory.""" - - t0 = time.monotonic() - tensors_to_pin = ( - [self._mmap_region._base] - if self._mmap_region is not None - else self.cpu_tensors - ) - num_pinned = 0 - for tensor in tensors_to_pin: - total_size_bytes = tensor.numel() * tensor.element_size() - result = torch.cuda.cudart().cudaHostRegister( - tensor.data_ptr(), total_size_bytes, 0 - ) - if result.value != 0: - logger.warning( - "cudaHostRegister failed for host tensor (code=%d) " - "- transfers will still work but may be slower (unpinned DMA)", - result.value, - ) - continue - if self._mmap_region is not None: - self._mmap_region.is_pinned = True - else: - self._manually_pinned_tensors.append(tensor) - num_pinned += 1 - - logger.debug( - "cudaHostRegister pin %.2f GB", - total_size_bytes / 1e9, - ) - - logger.info( - "Completed CPU memory pinning: %d tensors pinned in %.3f s", - num_pinned, - time.monotonic() - t0, - ) - def submit_store( self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec ) -> bool: From 07d33e575b472db52ae73ad44af18d909f34f177 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 29 Jun 2026 16:42:35 +0100 Subject: [PATCH 121/138] [MyPy] Fix mypy incompatible assignment errors in LRUCacheLoRAModelManager (#44657) Signed-off-by: Martin Hickey --- vllm/lora/model_manager.py | 46 ++++++-------------------------------- 1 file changed, 7 insertions(+), 39 deletions(-) diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index a24a75b8172..bc3d278af4e 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -98,14 +98,17 @@ class LoRAModelManager: f"No supported LoRA modules found in {self.model.__class__.__name__}." ) - self._registered_adapters: dict[int, LoRAModel] = {} - # Dict instead of a set for compatibility with LRUCache. - self._active_adapters: dict[int, None] = {} self.adapter_type = "LoRA" self.lora_config = lora_config self.device = device self.max_num_seqs = max_num_seqs assert self.capacity >= self.lora_slots + self._registered_adapters: AdapterLRUCache[LoRAModel] = AdapterLRUCache( + self.capacity, self.deactivate_adapter + ) + self._active_adapters: AdapterLRUCache[None] = AdapterLRUCache( + self.lora_slots, self._deactivate_adapter + ) self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8 self.lora_index_to_id: list[int | None] = [None] * self.lora_slots self.vocab_size = vocab_size @@ -1156,50 +1159,15 @@ class LoRAModelManager: return True def list_adapters(self) -> dict[int, LoRAModel]: - return dict(self._registered_adapters) + return dict(self._registered_adapters.cache) def get_adapter(self, adapter_id: int) -> LoRAModel | None: return self._registered_adapters.get(adapter_id) -class LoRALRUCache(AdapterLRUCache[LoRAModel]): - def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], object]): - super().__init__(capacity, deactivate_lora_fn) - - class LRUCacheLoRAModelManager(LoRAModelManager): """A model manager that manages multiple LoRAs with LRU cache.""" - def __init__( - self, - model: SupportsLoRAModel, - max_num_seqs: int, - max_num_batched_tokens: int, - vocab_size: int, - lora_config: LoRAConfig, - device: torch.device, - vllm_config: VllmConfig, - ): - super().__init__( - model, - max_num_seqs, - max_num_batched_tokens, - vocab_size, - lora_config, - device, - vllm_config, - ) - self._registered_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment] - self.capacity, self.deactivate_adapter - ) - self._active_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment] - self.lora_slots, self._deactivate_adapter - ) - - def list_adapters(self) -> dict[int, LoRAModel]: - """List all registered LoRAModels.""" - return dict(self._registered_adapters.cache) - def add_adapter(self, lora: LoRAModel) -> bool: """Add a LoRAModel to the manager.""" logger.debug("Adding lora. Model id: %d, int id: %d", lora.id, lora.id) From 379acd4e4fc33c3939556cf3a888f0963ec5c8ce Mon Sep 17 00:00:00 2001 From: HDCharles <39544797+HDCharles@users.noreply.github.com> Date: Mon, 29 Jun 2026 11:55:42 -0400 Subject: [PATCH 122/138] [Bugfix][Quantization] Fix W8A8 int-quantized scheme selection regression (#46860) Signed-off-by: HDCharles --- tests/quantization/test_compressed_tensors.py | 196 ++++++++++++++++++ .../compressed_tensors/compressed_tensors.py | 2 +- 2 files changed, 197 insertions(+), 1 deletion(-) diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index d51505a700a..626717cd4a3 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso CompressedTensorsW8A8Int8, CompressedTensorsW8A8Mxfp8, CompressedTensorsW8A16Fp8, + CompressedTensorsWNA8O8Int, CompressedTensorsWNA16, ) from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( @@ -672,6 +673,201 @@ def test_get_scheme_dict_returns_none_on_no_match(): assert result is None +# Test constants for activation quantization +_STATIC_SYM_INT8_ACT = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.TENSOR.value, + symmetric=True, + dynamic=False, +) + +_STATIC_ASYM_INT8_ACT = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.TENSOR.value, + symmetric=False, + dynamic=False, +) + +_DYNAMIC_INT8_ACT = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.TOKEN.value, + symmetric=True, + dynamic=True, +) + + +@pytest.mark.parametrize( + "weight_bits,weight_strategy,input_act,output_act,format,expected_scheme", + [ + # W8A8 int-quantized -> W8A8Int8 (regression test for #46389) + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _STATIC_SYM_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_channel_static_sym", + ), + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _STATIC_ASYM_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_channel_static_asym", + ), + pytest.param( + 8, + QuantizationStrategy.TENSOR.value, + _STATIC_SYM_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_tensor_static", + ), + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _DYNAMIC_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_channel_dynamic", + ), + # W8A8O8 int-quantized -> WNA8O8Int (both input and output) + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _STATIC_SYM_INT8_ACT, + _STATIC_SYM_INT8_ACT, + "int-quantized", + CompressedTensorsWNA8O8Int, + id="w8a8o8_channel", + ), + pytest.param( + 4, + QuantizationStrategy.GROUP.value, + _STATIC_SYM_INT8_ACT, + _STATIC_SYM_INT8_ACT, + "int-quantized", + CompressedTensorsWNA8O8Int, + id="w4a8o8_group", + ), + # Weight-only pack-quantized -> WNA16 + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w8_pack", + ), + pytest.param( + 4, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w4_pack", + ), + pytest.param( + 2, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w2_pack", + ), + pytest.param( + 3, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w3_pack", + ), + pytest.param( + 5, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w5_pack", + ), + pytest.param( + 6, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w6_pack", + ), + pytest.param( + 7, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w7_pack", + ), + ], +) +def test_scheme_selection( + weight_bits, weight_strategy, input_act, output_act, format, expected_scheme +): + """Test that _get_scheme_from_parts selects the correct scheme. + + This parametrized test verifies scheme selection for various combinations + of weight bits, quantization strategies, input/output activations, and + compression formats. + + Key regression test: W8A8 int-quantized models with channel-wise weights + should use W8A8Int8 (true int8 gemm), not WNA8O8Int (fake-quant). + WNA8O8Int should only match when BOTH input and output activations are + present. + """ + weight_quant = QuantizationArgs( + num_bits=weight_bits, + type=QuantizationType.INT, + strategy=weight_strategy, + symmetric=True, + dynamic=False, + group_size=128 if weight_strategy == QuantizationStrategy.GROUP.value else None, + ) + + config = CompressedTensorsConfig( + target_scheme_map={}, + ignore=[], + quant_format=format, + ) + + scheme = config._get_scheme_from_parts( + weight_quant=weight_quant, + input_quant=input_act, + output_quant=output_act, + format=format, + ) + + assert isinstance(scheme, expected_scheme), ( + f"Expected {expected_scheme.__name__} for " + f"W{weight_bits} {weight_strategy} + " + f"input_act={input_act} + output_act={output_act} + " + f"format={format}, got {type(scheme).__name__}" + ) + + @pytest.mark.skipif( not current_platform.is_cuda() or not current_platform.has_device_capability(75), reason="MXFP8 requires Turing (sm_75+) or newer.", diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index d52386d5d1a..2091a1cb6e4 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -679,7 +679,7 @@ class CompressedTensorsConfig(QuantizationConfig): and output_quant.num_bits == 8 and not output_quant.dynamic ) - return is_intN_weight and (is_static_int8_in or is_static_int8_out) + return is_intN_weight and (is_static_int8_in and is_static_int8_out) def _get_scheme_from_parts( self, From c8fb2963bd1baebbdd28062097096b59b2ba3189 Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Mon, 29 Jun 2026 12:28:32 -0400 Subject: [PATCH 123/138] [FS-Offloading] Batch Lookup in C (#46713) Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath --- CMakeLists.txt | 15 +++++ csrc/fs_io.cpp | 69 +++++++++++++++++++++ setup.py | 2 + tests/v1/kv_offload/tiering/test_fs_tier.py | 65 ++++++++++++++++++- vllm/v1/kv_offload/tiering/fs/manager.py | 13 +++- 5 files changed, 162 insertions(+), 2 deletions(-) create mode 100644 csrc/fs_io.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index cbd5583bbfd..1ef9d596aec 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -140,6 +140,21 @@ if(Python_VERSION VERSION_GREATER_EQUAL "3.11") WITH_SOABI) endif() +# +# fs_io extension (pure CXX; must stay above the non-CUDA device branch +# so CPU builds define the target before the early return). +# GIL-releasing filesystem helpers for FileSystemTierManager. +# +if(Python_VERSION VERSION_GREATER_EQUAL "3.11") + define_extension_target( + fs_io_C + DESTINATION vllm + LANGUAGE CXX + SOURCES csrc/fs_io.cpp + USE_SABI 3.11 + WITH_SOABI) +endif() + # # Forward the non-CUDA device extensions to external CMake scripts. # diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp new file mode 100644 index 00000000000..fdf3e614e64 --- /dev/null +++ b/csrc/fs_io.cpp @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include + +#include + +#include + +extern "C" { + +static void _batch_lookup(const std::vector& paths, + std::vector& exists_flags) { + for (size_t i = 0; i < paths.size(); i++) { + exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0; + } +} + +/// @brief Check file existence for a batch of paths. +/// @param paths list[str] โ€“ absolute paths to check. +/// @return list[bool] โ€“ True if the corresponding path exists, False otherwise. +/// @note Releases the GIL for the entire batch. File existence via access(2). +static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) { + PyObject* path_list; + if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(path_list); + std::vector paths(n); + for (Py_ssize_t i = 0; i < n; i++) { + paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr); + if (paths[i] == nullptr) { + return nullptr; + } + } + + std::vector exists_flags(n); + { + Py_BEGIN_ALLOW_THREADS _batch_lookup(paths, exists_flags); + Py_END_ALLOW_THREADS + } + + PyObject* result = PyList_New(n); + if (result == nullptr) { + return nullptr; + } + for (Py_ssize_t i = 0; i < n; i++) { + PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i])); + } + return result; +} + +static PyMethodDef fs_io_C_methods[] = { + {"batch_lookup", batch_lookup, METH_VARARGS, + "batch_lookup(paths: list[str]) -> list[bool]\n" + "\n" + "Check file existence for a batch of paths."}, + {nullptr, nullptr, 0, nullptr}, +}; + +static struct PyModuleDef fs_io_C_module = { + PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1, + fs_io_C_methods, +}; + +PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); } + +} // extern "C" diff --git a/setup.py b/setup.py index ad9aed07a31..b305fb1b00f 100644 --- a/setup.py +++ b/setup.py @@ -777,6 +777,7 @@ class precompiled_wheel_utils: "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", "vllm/cumem_allocator.abi3.so", "vllm/spinloop.abi3.so", + "vllm/fs_io_C.abi3.so", # ROCm-specific libraries "vllm/_rocm_C.abi3.so", } @@ -1104,6 +1105,7 @@ if _is_cuda() or _is_hip(): if sys.version_info >= (3, 11): ext_modules.append(CMakeExtension(name="vllm.spinloop")) + ext_modules.append(CMakeExtension(name="vllm.fs_io_C")) if _is_hip(): ext_modules.append(CMakeExtension(name="vllm._rocm_C")) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 7245ae1ba7a..0300fb5d4d4 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -103,7 +103,7 @@ def lookup_and_wait( keys: list[OffloadKey], ctx: ReqContext = _CTX, timeout: float = 1.0, -) -> list[bool]: +) -> list[LookupResult]: """Perform a full async lookup cycle and return resolved results.""" for k in keys: tier.lookup(k, ctx) @@ -332,3 +332,66 @@ def test_wait_idle_blocks_until_tasks_complete(): gate.set() pool.shutdown(wait=True) waiter.join(timeout=5.0) + + +def test_batch_lookup_c_extension(tmp_path): + """Validates batch_lookup_C: empty, single, all-existing, all-missing, + mixed ordering, and input type validation.""" + try: + from vllm.fs_io_C import batch_lookup as batch_lookup_C + except ImportError: + pytest.skip("fs_io_C extension not built") + + # Setup + all_exist = [str(tmp_path / f"e{i}.bin") for i in range(3)] + for p in all_exist: + open(p, "w").close() + all_missing = [str(tmp_path / f"m{i}.bin") for i in range(3)] + + # Empty list + assert batch_lookup_C([]) == [] + + # Single existing / missing + assert batch_lookup_C([all_exist[0]]) == [True] + assert batch_lookup_C([all_missing[0]]) == [False] + + # All existing / all missing + assert batch_lookup_C(all_exist) == [True, True, True] + assert batch_lookup_C(all_missing) == [False, False, False] + + # Mixed โ€” verifies index ordering is preserved + paths = [val for pair in zip(all_exist, all_missing) for val in pair] + assert batch_lookup_C(paths) == [True, False, True, False, True, False] + + # Input validation: non-list argument + with pytest.raises(TypeError): + batch_lookup_C(("/tmp/foo",)) + with pytest.raises(TypeError): + batch_lookup_C(None) + + # Input validation: non-str elements in list + with pytest.raises(TypeError): + batch_lookup_C([None]) + with pytest.raises(TypeError): + batch_lookup_C([b"/tmp/foo"]) + with pytest.raises(TypeError): + batch_lookup_C([42]) + with pytest.raises(TypeError): + batch_lookup_C([all_exist[0], None]) # valid first, invalid mid-list + + +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): + import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod + + if use_c_ext and not mgr_mod._HAS_BATCH_LOOKUP_C: + pytest.skip("fs_io_C extension not built") + + monkeypatch.setattr(mgr_mod, "_HAS_BATCH_LOOKUP_C", use_c_ext) + + tier, _ = fs_tier + tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(tier)) + + results = lookup_and_wait(tier, [key(1), key(2)]) + assert results == [LookupResult.HIT, LookupResult.MISS] diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 329a24daf34..816e88c5229 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -21,6 +21,13 @@ import os from collections.abc import Iterable from typing import TYPE_CHECKING +try: + from vllm.fs_io_C import batch_lookup as batch_lookup_C + + _HAS_BATCH_LOOKUP_C = True +except ImportError: + _HAS_BATCH_LOOKUP_C = False + from typing_extensions import override from vllm.logger import init_logger @@ -56,7 +63,11 @@ class FsAsyncLookupManager(AsyncLookupManager): def batch_lookup( self, keys: list[OffloadKey], req_context: ReqContext ) -> Iterable[bool]: - return (os.path.exists(self._tier.file_mapper.get_file_name(k)) for k in keys) + paths = [self._tier.file_mapper.get_file_name(k) for k in keys] + if _HAS_BATCH_LOOKUP_C: + # C extension: GIL released for the entire faccessat() batch. + return batch_lookup_C(paths) + return (os.path.exists(p) for p in paths) class FileSystemTierManager(SecondaryTierManager): From debec6440b89fe6ab14acb00e6eb2b04257f57a2 Mon Sep 17 00:00:00 2001 From: Jason Li Date: Mon, 29 Jun 2026 12:29:39 -0400 Subject: [PATCH 124/138] Add MiniMax-M3 modelopt nvfp4 support (#46756) Signed-off-by: Xin Li Signed-off-by: jasonlizhengjian Co-authored-by: Xin Li --- tests/quantization/test_modelopt.py | 6 ++ .../fused_moe/experts/trtllm_nvfp4_moe.py | 68 ++++++++++++++++--- .../layers/quantization/modelopt.py | 27 ++++++++ .../quantization/utils/flashinfer_utils.py | 1 + 4 files changed, 93 insertions(+), 9 deletions(-) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 0b54bcdbdfa..32450231487 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -18,6 +18,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, ModelOptMixedPrecisionConfig, + ModelOptMxFp8Config, ModelOptNvFp4Config, ModelOptNvFp4LinearMethod, ) @@ -84,6 +85,11 @@ def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionCon kv_cache_quant_algo=None, exclude_modules=[], ), + mxfp8_config=ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + ), ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index e45fc77ad90..518c87ce4df 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -66,16 +66,46 @@ class TrtLlmNvFp4ExpertsBase: else: self.g1_scale_c = self.quant_config.a2_gscale.clone() - if moe_config.is_act_and_mul and quant_config.gemm1_clamp_limit is not None: - device = torch.accelerator.current_device_index() - self.gemm1_clamp_limit = torch.full( + # Fall back to moe_config.swiglu_* when quant_config doesn't carry them + # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config). + device = torch.accelerator.current_device_index() + + def _per_expert(val: float | None) -> torch.Tensor | None: + if val is None: + return None + return torch.full( (self.local_num_experts,), - quant_config.gemm1_clamp_limit, + float(val), dtype=torch.float32, device=device, ) + + clamp = quant_config.gemm1_clamp_limit + if clamp is None: + clamp = getattr(moe_config, "swiglu_limit", None) + alpha = quant_config.gemm1_alpha + if alpha is None: + alpha = getattr(moe_config, "swiglu_alpha", None) + beta = quant_config.gemm1_beta + if beta is None: + beta = getattr(moe_config, "swiglu_beta", None) + + if moe_config.is_act_and_mul: + self.gemm1_clamp_limit = _per_expert(clamp) + self.gemm1_alpha = _per_expert(alpha) + self.gemm1_beta = _per_expert(beta) else: self.gemm1_clamp_limit = None + self.gemm1_alpha = None + self.gemm1_beta = None + + logger.debug_once( + "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", + moe_config.activation, + alpha, + beta, + clamp, + ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -109,6 +139,25 @@ class TrtLlmNvFp4ExpertsBase: ) self.gemm1_clamp_limit = layer.gemm1_clamp_limit + # beta shifts the raw GEMM1 accumulator, so fold by g1_alphas like the + # clamp limit. alpha is applied to the dequantized gate, so it stays + # raw. Register both on the layer so EPLB rearranges them with the + # other per-expert tensors. + if self.gemm1_beta is not None: + gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas + layer.register_parameter( + "gemm1_beta", + torch.nn.Parameter(gemm1_beta, requires_grad=False), + ) + self.gemm1_beta = layer.gemm1_beta + + if self.gemm1_alpha is not None: + layer.register_parameter( + "gemm1_alpha", + torch.nn.Parameter(self.gemm1_alpha, requires_grad=False), + ) + self.gemm1_alpha = layer.gemm1_alpha + @staticmethod def _supports_current_device() -> bool: """Supports only Blackwell-family GPUs.""" @@ -137,12 +186,13 @@ class TrtLlmNvFp4ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports only SiLU, RELU^2 non-gated and GELU activation.""" + """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI.""" return activation in [ MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU, MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, ] @staticmethod @@ -248,8 +298,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, gemm1_clamp_limit=self.gemm1_clamp_limit, gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), @@ -409,8 +459,8 @@ class TrtLlmNvFp4ExpertsMonolithic( gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, gemm1_clamp_limit=self.gemm1_clamp_limit, gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index d51a2dd312a..8fa1cb4d544 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2283,6 +2283,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): fp8_config: ModelOptFp8Config, nvfp4_config: ModelOptNvFp4Config, w4a16_nvfp4_config: ModelOptNvFp4Config, + mxfp8_config: ModelOptMxFp8Config, ) -> None: super().__init__(exclude_modules) self.kv_cache_quant_method = kv_cache_quant_method @@ -2290,6 +2291,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): self.fp8_config = fp8_config self.nvfp4_config = nvfp4_config self.w4a16_nvfp4_config = w4a16_nvfp4_config + self.mxfp8_config = mxfp8_config def get_name(self) -> QuantizationMethods: return "modelopt_mixed" @@ -2380,6 +2382,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): group_size=group_size, ) + mxfp8_config = ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo=kv_cache_quant_method, + exclude_modules=[], + ) + return cls( kv_cache_quant_method=kv_cache_quant_method, exclude_modules=exclude_modules, @@ -2387,6 +2395,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): fp8_config=fp8_config, nvfp4_config=nvfp4_config, w4a16_nvfp4_config=w4a16_nvfp4_config, + mxfp8_config=mxfp8_config, ) def _resolve_quant_algo(self, prefix: str) -> str | None: @@ -2441,6 +2450,17 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): if key.startswith(parent_dot): return info["quant_algo"].upper() + # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj). + for candidate in self._quantized_layer_prefix_candidates(prefix): + parent_dot = candidate.rsplit(".", 1)[0] + "." + algos = { + info["quant_algo"].upper() + for key, info in self.quantized_layers.items() + if key.startswith(parent_dot) and "." not in key[len(parent_dot) :] + } + if len(algos) == 1: + return algos.pop() + return None @staticmethod @@ -2486,6 +2506,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): return ModelOptNvFp4LinearMethod(self.nvfp4_config) if quant_algo == "W4A16_NVFP4": return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config) + if quant_algo == "MXFP8": + return ModelOptMxFp8LinearMethod(self.mxfp8_config) # Layer not in quantized_layers โ€” leave unquantized return UnquantizedLinearMethod() @@ -2505,6 +2527,11 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): quant_config=self.w4a16_nvfp4_config, moe_config=layer.moe_config, ) + if quant_algo == "MXFP8": + return ModelOptMxFp8FusedMoE( + quant_config=self.mxfp8_config, + moe_config=layer.moe_config, + ) return None return None diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 1cbfdf69c99..9961d0f0a12 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -36,6 +36,7 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType" MoEActivation.GELU: ActivationType.Geglu, MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, + MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, } return ACTIVATION_TO_FI_ACTIVATION[activation] From 4708292d48f3f15978a6ad3befe5f0052bc86491 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Mon, 29 Jun 2026 12:30:57 -0400 Subject: [PATCH 125/138] Bump flashinfer version to 0.6.13 (#46683) Signed-off-by: wzhao18 Co-authored-by: Jee Jee Li --- docker/Dockerfile | 2 +- docker/Dockerfile.nightly_torch | 4 ++-- docker/versions.json | 2 +- requirements/cuda.txt | 4 ++-- tests/evals/gsm8k/test_gsm8k_correctness.py | 11 +++++++++-- vllm/model_executor/warmup/kernel_warmup.py | 21 ++++++++++++++------- 6 files changed, 29 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index c86795586c3..945cb14bcb2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -793,7 +793,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.12 +ARG FLASHINFER_VERSION=0.6.13 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 0f2ec9f3a2e..cc706f59ae7 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -257,13 +257,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2. # build flashinfer for torch nightly from source around 10 mins -# release version: v0.6.12 +# release version: v0.6.13 # todo(elainewy): cache flashinfer build result for faster build ENV CCACHE_DIR=/root/.cache/ccache RUN --mount=type=cache,target=/root/.cache/ccache \ --mount=type=cache,target=/root/.cache/uv \ echo "git clone flashinfer..." \ - && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \ + && git clone --depth 1 --branch v0.6.13 --recursive https://github.com/flashinfer-ai/flashinfer.git \ && cd flashinfer \ && git submodule update --init --recursive \ && echo "finish git clone flashinfer..." \ diff --git a/docker/versions.json b/docker/versions.json index 3145cfcc53e..cfe1b1654f5 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.12" + "default": "0.6.13" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 124dae4846d..edaf9d2dd6a 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -10,8 +10,8 @@ torchaudio==2.11.0 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version PyNvVideoCodec==2.0.4 # FlashInfer should be updated together with the Dockerfile -flashinfer-python==0.6.12 -flashinfer-cubin==0.6.12 +flashinfer-python==0.6.13 +flashinfer-cubin==0.6.13 apache-tvm-ffi==0.1.9 tilelang==0.1.9 nvidia-cudnn-frontend>=1.19.1 diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index d14f41843b8..f796f910bb5 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -33,6 +33,8 @@ QUARK_MXFP4_TORCH_COMPATIBLE = find_spec("quark") is not None and ( else True ) +DEFAULT_STARTUP_MAX_WAIT_SECONDS = 1200 + def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict: """Run GSM8K evaluation using our isolated script.""" @@ -127,7 +129,11 @@ def test_gsm8k_correctness(config_filename): ] ) - env_dict = eval_config.get("env", None) + startup_max_wait_seconds = eval_config.get( + "startup_max_wait_seconds", DEFAULT_STARTUP_MAX_WAIT_SECONDS + ) + env_dict = dict(eval_config.get("env") or {}) + env_dict["VLLM_ENGINE_READY_TIMEOUT_S"] = str(int(startup_max_wait_seconds)) print(f"Starting GSM8K evaluation for model: {eval_config['model_name']}") print(f"Expected metric threshold: {eval_config['accuracy_threshold']}") @@ -139,6 +145,7 @@ def test_gsm8k_correctness(config_filename): "rocm_request_timeout_seconds", request_timeout_seconds ) print(f"Request timeout: {request_timeout_seconds}s") + print(f"Startup max wait: {startup_max_wait_seconds}s") print(f"Server args: {' '.join(server_args)}") print(f"Environment variables: {env_dict}") @@ -147,7 +154,7 @@ def test_gsm8k_correctness(config_filename): eval_config["model_name"], server_args, env_dict=env_dict, - max_wait_seconds=eval_config.get("startup_max_wait_seconds", 600), + max_wait_seconds=startup_max_wait_seconds, ) as remote_server: server_url = remote_server.url_for("v1") print(f"Server started at: {server_url}") diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 7edbff4d4a6..d9b71b2a0c7 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -113,12 +113,6 @@ def kernel_warmup(worker: "Worker"): ) -# TODO: remove once FlashInfer upstream fixes the persistent file cache -# to resolve collisions like `use_8x4_sf_layout=True/False`, which causes -# invalid tactics to be chosen -_FLASHINFER_USE_PERSISTENT_CACHE = False - - def flashinfer_autotune(runner: "GPUModelRunner") -> None: """ Autotune FlashInfer operations. @@ -135,7 +129,20 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: import vllm.utils.flashinfer as fi_utils from vllm.distributed.parallel_state import get_world_group - if not _FLASHINFER_USE_PERSISTENT_CACHE: + use_persistent_cache = True + + deepep_a2a_backends = { + "deepep_high_throughput", + "deepep_low_latency", + "deepep_v2", + } + if runner.vllm_config.parallel_config.all2all_backend in deepep_a2a_backends: + # DeepEP dispatch/combine can timeout when only rank 0 + # performs autotune and falls behind other ranks. + # Thus we skip persistent cache in this case. + use_persistent_cache = False + + if not use_persistent_cache: with torch.inference_mode(), fi_utils.autotune(): runner._dummy_run( num_tokens=runner.scheduler_config.max_num_batched_tokens, From 030c9523bdb6a6292545768c863fd747c195b06b Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Mon, 29 Jun 2026 18:40:34 +0200 Subject: [PATCH 126/138] [Perf][1/N] Expand Triton kernel warmup coverage, DSv4 (#46634) Signed-off-by: LopezCastroRoberto Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Lucas Wilkinson --- vllm/model_executor/warmup/kernel_warmup.py | 13 + .../warmup/sparse_mla_triton_warmup.py | 330 ++++++++++++++++++ .../warmup/v1_block_table_warmup.py | 43 +++ 3 files changed, 386 insertions(+) create mode 100644 vllm/model_executor/warmup/sparse_mla_triton_warmup.py create mode 100644 vllm/model_executor/warmup/v1_block_table_warmup.py diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d9b71b2a0c7..f1d7788a988 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -25,6 +25,12 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( flashinfer_sparse_mla_decode_autotune_warmup, ) from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup +from vllm.model_executor.warmup.sparse_mla_triton_warmup import ( + sparse_mla_triton_warmup_if_needed, +) +from vllm.model_executor.warmup.v1_block_table_warmup import ( + warm_v1_block_table_kernels, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -41,6 +47,12 @@ def kernel_warmup(worker: "Worker"): minimax_m3_msa_warmup, ) + # Pooling models do not use the generation slot-mapping path. + if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model: + warm_v1_block_table_kernels( + getattr(worker.model_runner, "device", torch.device("cuda")), + worker.scheduler_config.max_num_batched_tokens, + ) qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder @@ -55,6 +67,7 @@ def kernel_warmup(worker: "Worker"): ) # Run next so input-prep kernels JIT against pristine runner state. + sparse_mla_triton_warmup_if_needed(worker) flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) diff --git a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py new file mode 100644 index 00000000000..e932849face --- /dev/null +++ b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up sparse-MLA Triton metadata kernels.""" + +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset( + { + "FLASHMLA_SPARSE_DSV4", + "FLASHINFER_MLA_SPARSE_DSV4", + "ROCM_FLASHMLA_SPARSE_DSV4", + "DEEPSEEK_SPARSE_SWA", + } +) +_GENERIC_SPARSE_MLA_BACKENDS = frozenset( + { + "FLASHMLA_SPARSE", + "FLASHINFER_MLA_SPARSE", + "FLASHINFER_MLA_SPARSE_SM120", + } +) + +_SPARSE_PREFILL_METADATA_NUM_PREFILLS = (1, 2, 4, 8) +_SPARSE_PREFILL_METADATA_NUM_DECODES = (0, 1, 2) +_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS = (4, 128) +_PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS = (2, 3) +_PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS = ( + # query_slice_start offset, query_slice_stop offset + (0, 0), + (0, -1), + (1, 0), + (1, -1), +) +_COMBINE_TOPK_SWA_INPUT_VARIANTS = ( + # offset_topk, offset_query_and_seq, offset_gather + (False, False, False), + (False, True, False), + (True, True, True), +) +_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES = ( + # compress_ratio, topk, topk_width, N + (1, 0, 512, 512), + (4, 512, 512, 512 * 4), + # DSv4-Pro C4A traffic uses top-k 1024 with N=1024. + (4, 1024, 1024, 1024), + (128, 8192, 8192, 8192 * 128), + # Real C128A traffic also specializes N=1 in one call path. + (128, 8192, 8192, 1), +) + + +def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int: + return max(0, min(num_tokens, max_tokens)) + + +def _next_power_of_2(x: int) -> int: + return 1 << (x - 1).bit_length() + + +def _hf_config_int(runner: "GPUModelRunner", name: str, default: int) -> int: + model_config = getattr(runner.vllm_config, "model_config", None) + hf_config = getattr(model_config, "hf_config", None) + return int(getattr(hf_config, name, default) or default) + + +def _attention_backend_name(backend: object) -> str | None: + get_name = getattr(backend, "get_name", None) + if get_name is None: + return None + try: + return get_name() + except NotImplementedError: + return None + + +def _has_attention_backend( + runner: "GPUModelRunner", + backend_names: frozenset[str], +) -> bool: + for groups in getattr(runner, "attn_groups", []) or (): + for group in groups: + name = _attention_backend_name(getattr(group, "backend", None)) + if name in backend_names: + return True + return False + + +def _warm_sparse_swa_prefill_metadata_kernel( + device: torch.device, + window_size: int, + prefill_tokens: int, +) -> None: + from vllm.v1.attention.backends.mla.sparse_swa import ( + _compute_prefill_metadata_kernel, + ) + + for num_prefills in _SPARSE_PREFILL_METADATA_NUM_PREFILLS: + for num_decodes in _SPARSE_PREFILL_METADATA_NUM_DECODES: + query_lens = [1] * num_decodes + query_lens += [prefill_tokens] * num_prefills + query_start_locs = [0] + for query_len in query_lens: + query_start_locs.append(query_start_locs[-1] + query_len) + query_start_loc = torch.tensor( + query_start_locs, + dtype=torch.int32, + device=device, + ) + seq_lens = torch.tensor( + [1] * num_decodes + [window_size + q for q in query_lens[num_decodes:]], + dtype=torch.int32, + device=device, + ) + prefill_gather_lens = torch.empty( + num_prefills, dtype=torch.int32, device=device + ) + _compute_prefill_metadata_kernel[(1,)]( + prefill_gather_lens, + seq_lens, + query_start_loc, + num_prefills, + num_decodes, + window_size, + BLOCK_SIZE=_next_power_of_2(num_prefills), + ) + + +def _warm_prefill_chunk_metadata_kernel( + device: torch.device, + compress_ratio: int, + query_len: int, +) -> None: + from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata + + num_reqs = 2 + query_start_loc_cpu = torch.arange( + 0, (num_reqs + 1) * query_len, query_len, dtype=torch.int32 + ) + query_start_loc = query_start_loc_cpu.to(device=device) + + uncompressed_seq_lens_cpu = torch.tensor( + [ + compress_ratio * multiplier + query_len + for multiplier in _PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS + ], + dtype=torch.int32, + ) + compressed_seq_lens_cpu = uncompressed_seq_lens_cpu // compress_ratio + uncompressed_seq_lens = uncompressed_seq_lens_cpu.to(device=device) + compressed_seq_lens = compressed_seq_lens_cpu.to(device=device) + block_table = torch.zeros( + (num_reqs, int(compressed_seq_lens_cpu.max().item())), + dtype=torch.int32, + device=device, + ) + + offset_uncompressed_seq_lens = torch.empty( + num_reqs + 1, dtype=torch.int32, device=device + )[1:] + offset_uncompressed_seq_lens.copy_(uncompressed_seq_lens) + query_slices = tuple( + slice(start, num_reqs * query_len + stop) + for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS + ) + for warmup_uncompressed_seq_lens in ( + uncompressed_seq_lens, + offset_uncompressed_seq_lens, + ): + for query_slice in query_slices: + build_prefill_chunk_metadata( + 0, + num_reqs, + query_start_loc, + query_start_loc_cpu, + warmup_uncompressed_seq_lens, + compressed_seq_lens, + compressed_seq_lens_cpu, + block_table, + compress_ratio, + query_slice=query_slice, + ) + + +def _warm_combine_topk_swa_indices_kernel( + device: torch.device, + num_tokens: int, + window_size: int, + compress_ratio: int, + topk: int, + topk_width: int, + n: int, +) -> None: + from vllm.models.deepseek_v4.common.ops.cache_utils import combine_topk_swa_indices + + if num_tokens <= 0: + return + + def _make_topk_indices(*, offset: bool) -> torch.Tensor: + if offset: + topk_storage = torch.full( + (num_tokens * topk_width + 1,), + -1, + dtype=torch.int32, + device=device, + ) + topk_indices = topk_storage[1:].reshape(num_tokens, topk_width) + else: + topk_indices = torch.full( + (num_tokens, topk_width), -1, dtype=torch.int32, device=device + ) + if topk > 0: + topk_indices.copy_( + torch.arange(num_tokens * topk_width, dtype=torch.int32, device=device) + .reshape(num_tokens, topk_width) + .remainder(topk_width) + ) + return topk_indices + + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + seq_lens = torch.tensor( + [window_size + num_tokens], dtype=torch.int32, device=device + ) + gather_lens = torch.tensor( + [min(window_size + num_tokens, window_size + num_tokens - 1)], + dtype=torch.int32, + device=device, + ) + offset_query_start_loc = torch.empty(3, dtype=torch.int32, device=device)[1:] + offset_query_start_loc.copy_(query_start_loc) + offset_seq_lens = torch.empty(2, dtype=torch.int32, device=device)[1:] + offset_seq_lens.copy_(seq_lens) + offset_gather_lens = torch.empty(2, dtype=torch.int32, device=device)[1:] + offset_gather_lens.copy_(gather_lens) + + for ( + offset_topk, + offset_query_and_seq, + offset_gather, + ) in _COMBINE_TOPK_SWA_INPUT_VARIANTS: + warmup_topk_indices = _make_topk_indices(offset=offset_topk) + warmup_query_start_loc = ( + offset_query_start_loc if offset_query_and_seq else query_start_loc + ) + warmup_seq_lens = offset_seq_lens if offset_query_and_seq else seq_lens + warmup_gather_lens = offset_gather_lens if offset_gather else gather_lens + n_values = (n,) if n == 1 else (n, n + 1) + for m in (window_size + num_tokens, topk_width): + for n_value in n_values: + combine_topk_swa_indices( + warmup_topk_indices, + warmup_query_start_loc, + warmup_seq_lens, + warmup_gather_lens, + window_size, + compress_ratio, + topk, + M=m, + N=n_value, + ) + + +@torch.inference_mode() +def sparse_mla_triton_warmup( + runner: "GPUModelRunner", + num_tokens: int, + *, + compress_ratios: tuple[int, ...], + combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (), +) -> None: + device = getattr(runner, "device", torch.device("cuda")) + window_size = _hf_config_int(runner, "sliding_window", 128) + + _warm_sparse_swa_prefill_metadata_kernel(device, window_size, num_tokens) + for compress_ratio in compress_ratios: + _warm_prefill_chunk_metadata_kernel(device, compress_ratio, num_tokens) + for compress_ratio, topk, topk_width, n in combine_topk_swa_cases: + _warm_combine_topk_swa_indices_kernel( + device, + num_tokens, + window_size, + compress_ratio, + topk, + topk_width, + n, + ) + + +def deepseek_v4_sparse_triton_warmup( + runner: "GPUModelRunner", + num_tokens: int, +) -> None: + sparse_mla_triton_warmup( + runner, + num_tokens, + compress_ratios=_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS, + combine_topk_swa_cases=_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES, + ) + + +def sparse_mla_triton_warmup_if_needed(worker: "Worker") -> None: + runner = worker.model_runner + if runner.is_pooling_model: + return + + max_tokens = worker.scheduler_config.max_num_batched_tokens + num_tokens = _clamp_warmup_tokens(8, max_tokens) + if num_tokens <= 0: + return + + try: + if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS): + deepseek_v4_sparse_triton_warmup(runner, num_tokens) + elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS): + sparse_mla_triton_warmup( + runner, + num_tokens, + compress_ratios=(1,), + ) + except Exception: + logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True) diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py new file mode 100644 index 00000000000..8d2328432eb --- /dev/null +++ b/vllm/model_executor/warmup/v1_block_table_warmup.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up v1 block-table Triton kernels.""" + +import torch + +_SLOT_MAPPING_WARMUP_TOKENS = 8 +_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16) +_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1 + + +def warm_v1_block_table_kernels( + device: torch.device, + max_tokens: int, +) -> None: + from vllm.v1.worker.block_table import BlockTable + + num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens)) + if num_tokens <= 0: + return + + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES: + max_num_blocks_per_req = max( + 1, (max(num_tokens, max_tokens) + block_size - 1) // block_size + ) + max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16 + block_table = BlockTable( + block_size=block_size, + max_num_reqs=1, + max_num_blocks_per_req=max_num_blocks_per_req, + max_num_batched_tokens=max(num_tokens, max_tokens), + pin_memory=False, + device=device, + kernel_block_size=block_size, + cp_kv_cache_interleave_size=( + _SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE + ), + ) + block_table.add_row(list(range(max_num_blocks_per_req)), 0) + block_table.commit_block_table(1) + block_table.compute_slot_mapping(1, query_start_loc, positions) From 7be582697b27277e2756a3878f563fa9dfea30aa Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Tue, 30 Jun 2026 00:44:05 +0800 Subject: [PATCH 127/138] [Bugfix] Fix DeepseekV2Model hidden_size (#46986) Signed-off-by: Jee Jee Li --- vllm/model_executor/models/deepseek_v2.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 144ff3971a2..211c88129c8 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1336,7 +1336,7 @@ class DeepseekV2Model(nn.Module): quant_config = vllm_config.quant_config self.config = config self.device = current_platform.device_type - + self.hidden_size = config.hidden_size self.vocab_size = config.vocab_size self.is_v32 = hasattr(config, "index_topk") if self.is_v32: @@ -1353,7 +1353,7 @@ class DeepseekV2Model(nn.Module): if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, - config.hidden_size, + self.hidden_size, quant_config=quant_config, prefix=f"{prefix}.embed_tokens", ) @@ -1370,11 +1370,11 @@ class DeepseekV2Model(nn.Module): ) if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps) else: self.norm = PPMissingLayer() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size + ["hidden_states", "residual"], self.hidden_size ) self.aux_hidden_state_layers = tuple[int, ...]() From 8ad4a01825ef941e785b2bc305ac7a6b5ca9c530 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 29 Jun 2026 17:56:17 +0100 Subject: [PATCH 128/138] [ModelRunner V2] Simplify recent UnlimitedOCR-related changes (#46975) Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/input_batch.py | 5 +++-- vllm/v1/worker/gpu/model_runner.py | 21 +++++-------------- vllm/v1/worker/gpu/model_states/default.py | 2 +- .../gpu/model_states/encoder_decoder.py | 2 +- .../worker/gpu/model_states/mamba_hybrid.py | 2 +- 5 files changed, 11 insertions(+), 21 deletions(-) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index a6a2b296e38..006e11e4500 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -96,8 +96,8 @@ class InputBatch: # Whether any requests in batch use structured output. has_structured_output_reqs: bool - # [num_reqs_after_padding] per-request prompt length for R-SWA (optional). - rswa_prefix_lens: torch.Tensor | None = None + # [num_reqs] per-request prompt length, only populated for R-SWA. + prompt_lens: torch.Tensor | None @classmethod def make_dummy( @@ -178,6 +178,7 @@ class InputBatch: cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=False, + prompt_lens=None, ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0f57e8a31cd..8869e93f9fa 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -223,13 +223,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_num_tokens=self.max_num_tokens, device=self.device, ) - # R-SWA: persistent GPU buffer for per-request prefix lengths (CUDA-graph safe). - self.rswa_prefix_lens_buffer: torch.Tensor | None = None - if self.model_config.rswa_window is not None: - self.rswa_prefix_lens_buffer = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=self.device - ) - if self.use_pp: self.pp_handler = PPHandler( max_num_reqs=self.max_num_reqs, @@ -992,14 +985,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): # max_seq_len is only consumed by the PP `compute_need_sampled_mask` max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np] - rswa_prefix_lens = None - if self.rswa_prefix_lens_buffer is not None: - rswa_prefix_lens = self.rswa_prefix_lens_buffer[:num_reqs_padded] - rswa_prefix_lens[:num_reqs] = self.req_states.prompt_len.gpu[ - idx_mapping[:num_reqs] - ] - if num_reqs_padded > num_reqs: - rswa_prefix_lens[num_reqs:].zero_() + prompt_lens = None + if self.model_config.rswa_window is not None: + # prompt_lens is only used in R-SWA case. + prompt_lens = self.req_states.prompt_len.gpu[idx_mapping] return InputBatch( req_ids=req_ids, @@ -1031,7 +1020,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=scheduler_output.has_structured_output_requests, - rswa_prefix_lens=rswa_prefix_lens, + prompt_lens=prompt_lens, ) def prepare_attn( diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index f760fc36dea..e5e89da2b2e 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -168,6 +168,6 @@ class DefaultModelState(ModelState): dcp_local_seq_lens=input_batch.dcp_local_seq_lens, positions=input_batch.positions, for_cudagraph_capture=for_capture, - rswa_prefix_lens=input_batch.rswa_prefix_lens, + rswa_prefix_lens=input_batch.prompt_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index 9edda27538e..f759c0b1e15 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -146,7 +146,7 @@ class EncoderDecoderModelState(ModelState): dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=enc_dec_attn_metadata, for_cudagraph_capture=for_capture, - rswa_prefix_lens=input_batch.rswa_prefix_lens, + rswa_prefix_lens=input_batch.prompt_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index e08b09f1895..a0f5968361c 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -141,7 +141,7 @@ class MambaHybridModelState(DefaultModelState): dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=mamba_attn_metadata, for_cudagraph_capture=for_capture, - rswa_prefix_lens=input_batch.rswa_prefix_lens, + rswa_prefix_lens=input_batch.prompt_lens, ) def postprocess_state( From 72f639927f6caf3495d69dec63b9d4a87ed782ef Mon Sep 17 00:00:00 2001 From: zofia <110436990+zufangzhu@users.noreply.github.com> Date: Tue, 30 Jun 2026 03:03:06 +0800 Subject: [PATCH 129/138] [XPU] [RMSNorm] revert weightless change on xpu (#46987) Signed-off-by: Zhu, Zufang Co-authored-by: Kunshang Ji --- vllm/kernels/xpu_ops.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/vllm/kernels/xpu_ops.py b/vllm/kernels/xpu_ops.py index 8a86b1226b4..df82962d802 100644 --- a/vllm/kernels/xpu_ops.py +++ b/vllm/kernels/xpu_ops.py @@ -31,10 +31,8 @@ def rms_norm( ) -> Tensor: assert variance_size is None if weight is None: - # Weightless _C ops are CUDA-only; native skips the multiply on XPU. - return ir.ops.rms_norm.impls["native"].impl_fn( - x, weight, epsilon, variance_size - ) + # Kernel requires weight tensor, pass ones + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) output = torch.empty(x.shape, device=x.device, dtype=x.dtype) torch.ops._C.rms_norm(output, x, weight, epsilon) return output @@ -61,12 +59,7 @@ def fused_add_rms_norm( ) -> tuple[Tensor, Tensor]: assert variance_size is None if weight is None: - # Weightless _C ops are CUDA-only; native skips the multiply on XPU. - output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn( - x, x_residual, weight, epsilon, variance_size - ) - x.copy_(output) - x_residual.copy_(residual) - return x, x_residual + # Kernel requires weight tensor, pass ones + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) return x, x_residual From a309d4fe60bef7657b88805a1fcc9b014c414314 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 29 Jun 2026 13:24:29 -0700 Subject: [PATCH 130/138] Support DCP with FlashInfer MLA (#43729) Signed-off-by: Woosuk Kwon --- docs/design/attention_backends.md | 2 +- vllm/v1/attention/backends/mla/flashinfer_mla.py | 14 ++++++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 9278ab6761a..f3067dfc859 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -220,7 +220,7 @@ MLA decode backends are selected using the standard | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | โŒ | โŒ | โŒ | โŒ | โœ… | Decoder | 10.x | -| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | โŒ | โŒ | โŒ | โŒ | โŒ | Decoder | 10.x | +| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | โŒ | โŒ | โŒ | โŒ | โœ… | Decoder | 10.x | | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | โŒ | โŒ | โŒ | โŒ | โŒ | Decoder | 10.x | | `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | โŒ | โŒ | โŒ | โŒ | โŒ | Decoder | 12.x | | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | โŒ | โŒ | โŒ | โŒ | โœ… | Decoder | 9.x-10.x | diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index 25ab3d7f659..07e2e44140e 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -113,6 +113,8 @@ g_fi_workspace = torch.zeros( class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): + can_return_lse_for_decode: bool = True + def __init__( self, num_heads: int, @@ -196,7 +198,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): if is_quantized_kv_cache(self.kv_cache_dtype): self.bmm2_scale *= layer._k_scale_float - o = trtllm_batch_decode_with_kv_cache_mla( + return_lse = self.need_to_return_lse_for_decode + kernel_out = trtllm_batch_decode_with_kv_cache_mla( query=q, kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), workspace_buffer=self._workspace_buffer, @@ -208,11 +211,14 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): max_seq_len=attn_metadata.max_seq_len, bmm1_scale=self.bmm1_scale, bmm2_scale=self.bmm2_scale, + return_lse=return_lse, ) + if return_lse: + o, lse = kernel_out + else: + o, lse = kernel_out, None # Flatten the output for consistent shape o = o.view(-1, o.shape[-2], o.shape[-1]) - # TODO: Return LSE pending support from Flashinfer API: - # https://github.com/flashinfer-ai/flashinfer/pull/1566 - return o, None + return o, lse From 61ab70ec3bd13dd422b86f3b80207d322994a5e7 Mon Sep 17 00:00:00 2001 From: zhrrr <43847754+izhuhaoran@users.noreply.github.com> Date: Tue, 30 Jun 2026 05:09:16 +0800 Subject: [PATCH 131/138] [Model Runner V2] support mamba hybrid models align prefix cache (#42406) Signed-off-by: zhuhaoran --- .../kernels/mamba/test_precopy_mamba_align.py | 180 +++++++ .../v1/e2e/general/test_mamba_prefix_cache.py | 296 ++++++++++- vllm/config/vllm.py | 11 - vllm/model_executor/models/diffusion_gemma.py | 4 +- vllm/v1/worker/gpu/model_runner.py | 14 +- vllm/v1/worker/gpu/model_states/interface.py | 17 +- .../worker/gpu/model_states/mamba_hybrid.py | 182 ++++++- vllm/v1/worker/gpu/warmup.py | 15 +- vllm/v1/worker/mamba_utils.py | 460 ++++++++++++++---- 9 files changed, 1046 insertions(+), 133 deletions(-) create mode 100644 tests/kernels/mamba/test_precopy_mamba_align.py diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py new file mode 100644 index 00000000000..be1e4559486 --- /dev/null +++ b/tests/kernels/mamba/test_precopy_mamba_align.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Equivalence test for ``precopy_mamba_align_fused_kernel``. + +The V2 "align" pre-copy must migrate mamba state across block boundaries with +byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` / +``get_temporal_copy_spec``): + +* conv state (SD layout, conv_width > 0): shift the sliding window by + ``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` -> + ``state[bt[dst_col], :conv_width - token_bias]``. +* temporal state (conv_width == 0): ``token_bias`` selects the accepted + speculative column -- ``state[bt[src_col + token_bias]]`` -> + ``state[bt[dst_col]]``. + +The kernel must also no-op when ``src_col < 0`` (fresh request) or +``src_col == dst_col`` (no boundary crossed). +""" + +from __future__ import annotations + +import torch + +from vllm.platforms import current_platform +from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel + +try: + import pytest + + pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), + reason="precopy_mamba_align_fused_kernel needs CUDA/Triton", + ) + _parametrize = pytest.mark.parametrize +except ModuleNotFoundError: # allow running directly as ``python `` + pytest = None + + def _parametrize(_name, _values): + def _deco(fn): + return fn + + return _deco + + +NUM_LAYERS = 3 +CONV_WIDTH = 4 # conv_kernel - 1 + num_spec +CONV_DIM = 96 +SSM_SHAPE = (4, 16, 16) +MAX_COLS = 8 + + +def _build_state(num_blocks, device): + """Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools.""" + convs, ssms = [], [] + for _ in range(NUM_LAYERS): + convs.append( + torch.randn( + num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device + ) + ) + ssms.append( + torch.randn(num_blocks, *SSM_SHAPE, dtype=torch.float32, device=device) + ) + return convs, ssms + + +def _build_meta(convs, ssms, device): + """Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer.""" + n = NUM_LAYERS * 2 + base = torch.zeros(n, dtype=torch.int64, device=device) + blk_stride = torch.zeros(n, dtype=torch.int64, device=device) + elem = torch.zeros(n, dtype=torch.int32, device=device) + inner = torch.zeros(n, dtype=torch.int64, device=device) + width = torch.zeros(n, dtype=torch.int32, device=device) + group = torch.zeros(n, dtype=torch.int32, device=device) + drc = torch.zeros(n, dtype=torch.int32, device=device) # DS rows (unused, SD) + drs = torch.zeros(n, dtype=torch.int64, device=device) + i = 0 + for layer in range(NUM_LAYERS): + conv, ssm = convs[layer], ssms[layer] + # conv (SD): width = size(1), inner = stride(1) + base[i] = conv.data_ptr() + blk_stride[i] = conv.stride(0) * conv.element_size() + elem[i] = conv.element_size() + width[i] = conv.size(1) + inner[i] = conv.stride(1) + i += 1 + # ssm (temporal): width = 0, inner = elems per block + base[i] = ssm.data_ptr() + blk_stride[i] = ssm.stride(0) * ssm.element_size() + elem[i] = ssm.element_size() + width[i] = 0 + inner[i] = ssm[0].numel() + i += 1 + return base, blk_stride, elem, inner, width, group, drc, drs + + +def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs): + """Apply the V1 copy semantics on clones, reading from the pre-copy state.""" + conv_pre = [c.clone() for c in convs] + ssm_pre = [s.clone() for s in ssms] + conv_ref = [c.clone() for c in convs] + ssm_ref = [s.clone() for s in ssms] + for r in range(num_reqs): + sc, dc, tb = int(src_col[r]), int(dst_col[r]), int(bias[r]) + if sc < 0 or sc == dc: + continue + sblk, dblk = int(bt[r, sc]), int(bt[r, dc]) + tblk = int(bt[r, sc + tb]) # temporal src column shifted by bias + for layer in range(NUM_LAYERS): + conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:] + ssm_ref[layer][dblk] = ssm_pre[layer][tblk] + return conv_ref, ssm_ref + + +@_parametrize("num_reqs", [1, 4, 16]) +@_parametrize("token_bias", [0, 1, 2]) +def test_precopy_matches_v1_copy_specs(num_reqs, token_bias): + device = torch.device("cuda") + torch.manual_seed(0) + # Distinct physical block per (req, col) so copies never alias. + num_blocks = num_reqs * MAX_COLS + 1 + bt = torch.empty(num_reqs, MAX_COLS, dtype=torch.int32, device=device) + for r in range(num_reqs): + bt[r] = torch.arange( + 1 + r * MAX_COLS, 1 + (r + 1) * MAX_COLS, dtype=torch.int32, device=device + ) + + # Per-req columns: req 0 fresh (src=-1, skip), req 1 same block (skip), + # the rest cross from col 1 -> col 0 with the given spec token bias. + src_col = torch.full((num_reqs,), 1, dtype=torch.int32, device=device) + dst_col = torch.zeros(num_reqs, dtype=torch.int32, device=device) + bias = torch.full((num_reqs,), token_bias, dtype=torch.int32, device=device) + if num_reqs >= 1: + src_col[0] = -1 # fresh -> no copy + if num_reqs >= 2: + dst_col[1] = 1 # src_col == dst_col -> no copy + + convs, ssms = _build_state(num_blocks, device) + conv_ref, ssm_ref = _reference( + convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs + ) + + base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( + convs, ssms, device + ) + bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + grid = (num_reqs, NUM_LAYERS * 2) + precopy_mamba_align_fused_kernel[grid]( + dst_col, + src_col, + bias, + bt_ptrs, + bt.stride(0), + base, + blk_stride, + elem, + inner, + width, + group, + drc, + drs, + idx_mapping, + num_reqs, + COPY_BLOCK_SIZE=1024, + CONV_STATE_DIM_FIRST=False, + ) + torch.accelerator.synchronize() + + for layer in range(NUM_LAYERS): + torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0) + torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) + + +if __name__ == "__main__": + for nr in (1, 4, 16): + for tb in (0, 1, 2): + test_precopy_matches_v1_copy_specs(nr, tb) + print(f"OK num_reqs={nr} token_bias={tb}") diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index e857b127285..4644a6cc7e1 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -11,6 +11,7 @@ import datasets import pytest import torch +import vllm.envs as envs from tests.utils import create_new_process_for_each_test from vllm import LLM, SamplingParams, TokensPrompt from vllm.config import CacheConfig @@ -494,12 +495,7 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn) -@create_new_process_for_each_test() -def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch): - run_ref_mamba_state_in_subprocess() - apply_patch(monkeypatch) - prompt_dataset = datasets.load_dataset("heheda/a_long_article") - full_prompt = prompt_dataset["train"][0]["text"] +def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: tests = { "accept_1": TestConfig( num_prompt_tokens=554, @@ -731,6 +727,27 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch): ), } + return tests + + +def fill_following_kv_cache_block_ids(test_config: TestConfig) -> None: + for step_action_prev, step_action_next in zip( + test_config.step_actions[:-1], test_config.step_actions[1:] + ): + if len(step_action_next.kv_cache_block_ids) == 0: + step_action_next.kv_cache_block_ids = ( + step_action_prev.kv_cache_block_ids.copy() + ) + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): + run_ref_mamba_state_in_subprocess() + apply_patch(monkeypatch) + prompt_dataset = datasets.load_dataset("heheda/a_long_article") + full_prompt = prompt_dataset["train"][0]["text"] + tests = get_mamba_prefix_cache_step_configs() + engine = LLM( model=MODEL, enable_prefix_caching=True, @@ -758,16 +775,7 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch): ) global cur_step_action_idx cur_step_action_idx = 0 - for step_action_prev, step_action_next in zip( - test_config.step_actions[:-1], test_config.step_actions[1:] - ): - if ( - step_action_next.kv_cache_block_ids is not None - and len(step_action_next.kv_cache_block_ids) == 0 - ): - prev_block_ids = step_action_prev.kv_cache_block_ids - if prev_block_ids is not None: - step_action_next.kv_cache_block_ids = prev_block_ids.copy() + fill_following_kv_cache_block_ids(test_config) global step_actions step_actions = test_config.step_actions _ = engine.generate( @@ -787,3 +795,259 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch): del engine torch.accelerator.empty_cache() cleanup_dist_env_and_memory() + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + envs.disable_envs_cache() + + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as MRV2GPUModelRunner + from vllm.v1.worker.gpu.model_states.mamba_hybrid import ( + MambaHybridModelState, + ) + from vllm.v1.worker.gpu.sample.output import SamplerOutput as MRV2SamplerOutput + + events: list[int] = [] + original_execute_model = MRV2GPUModelRunner.execute_model + original_sample = MRV2GPUModelRunner.sample + original_preprocess_state = MambaHybridModelState.preprocess_state + original_postprocess_state = MambaHybridModelState.postprocess_state + original_step_action_fn = InprocClient.get_output + original_allocate_slots = KVCacheManager.allocate_slots + captured: dict[str, Any] = {} + + def temporal_states(model_state, block_tables, kv_cache_config): + # Qwen3-Next keeps the temporal (ssm) state as the last Mamba cache. + forward_context = ( + model_state.vllm_config.compilation_config.static_forward_context + ) + group_ids, _ = get_mamba_groups(kv_cache_config) + for group_id in group_ids: + block_table = block_tables[group_id] + for layer_name in kv_cache_config.kv_cache_groups[group_id].layer_names: + yield forward_context[layer_name].kv_cache[-1], block_table + + def temporal_block(temporal_state, block_table, col): + return temporal_state[int(block_table[0, col].item())] + + def wrapped_preprocess_state( + self: MambaHybridModelState, + input_batch: Any, + block_tables: tuple[torch.Tensor, ...], + kv_cache_config: KVCacheConfig, + num_computed_tokens: torch.Tensor, + ) -> None: + captured["block_tables"] = block_tables + captured["kv_cache_config"] = kv_cache_config + expected = ( + None if cur_step_action is None else cur_step_action.preprocess_copy_idx + ) + snapshots = [] + if expected is not None and expected != (-1, -1): + for temporal, bt in temporal_states(self, block_tables, kv_cache_config): + snapshots.append( + (temporal, bt, temporal_block(temporal, bt, expected[0]).clone()) + ) + ret = original_preprocess_state( + self, input_batch, block_tables, kv_cache_config, num_computed_tokens + ) + if cur_step_action is not None: + req_idx = int(input_batch.idx_mapping[0].item()) + src_col = int(self._mamba_src_col_gpu[req_idx].item()) + off = int(self._mamba_src_off_gpu[req_idx].item()) + dst = int(self._mamba_state_idx_gpu[req_idx].item()) + actual = (-1, -1) if src_col < 0 or src_col == dst else (src_col + off, dst) + assert actual == expected, ( + f"V2 align preprocess copy: expected={expected}, " + f"actual={actual}, {cur_step_action=}" + ) + for temporal, bt, src_state in snapshots: + torch.testing.assert_close( + temporal_block(temporal, bt, expected[1]), src_state + ) + return ret + + def wrapped_postprocess_state( + self: MambaHybridModelState, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor | int, + num_computed_tokens: torch.Tensor | None = None, + ) -> None: + action = cur_step_action + block_tables = captured.get("block_tables") + kv_cache_config = captured.get("kv_cache_config") + # The postprocess kernel does not expose its indices, so only the copy + # case is checked, by effect: snapshot the src block, expect dst == src. + if ( + action is None + or num_computed_tokens is None + or block_tables is None + or action.postprocess_copy_idx == (-1, -1) + ): + return original_postprocess_state( + self, idx_mapping, num_sampled, num_computed_tokens + ) + expected = action.postprocess_copy_idx + snapshots = [ + (temporal, bt, temporal_block(temporal, bt, expected[0]).clone()) + for temporal, bt in temporal_states(self, block_tables, kv_cache_config) + ] + ret = original_postprocess_state( + self, idx_mapping, num_sampled, num_computed_tokens + ) + for temporal, bt, src_state in snapshots: + torch.testing.assert_close( + temporal_block(temporal, bt, expected[1]), src_state + ) + return ret + + def wrapped_execute_model( + self: MRV2GPUModelRunner, + scheduler_output: SchedulerOutput, + *args: Any, + **kwargs: Any, + ): + events.extend( + req.num_computed_tokens for req in scheduler_output.scheduled_new_reqs + ) + events.extend(scheduler_output.scheduled_cached_reqs.num_computed_tokens) + if cur_step_action is not None: + num_scheduled_tokens = next( + iter(scheduler_output.num_scheduled_tokens.values()) + ) + assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens + ret = original_execute_model(self, scheduler_output, *args, **kwargs) + if cur_step_action is not None and self.execute_model_state is not None: + input_batch = self.execute_model_state.input_batch + assert ( + cur_step_action.num_computed_tokens_start + == input_batch.positions[input_batch.query_start_loc[0]].item() + ) + return ret + + def fake_sample( + self: MRV2GPUModelRunner, + hidden_states: torch.Tensor, + input_batch: Any, + grammar_output: Any, + ): + if cur_step_action is None: + return original_sample(self, hidden_states, input_batch, grammar_output) + + num_reqs = input_batch.num_reqs + sampled_token_ids = torch.ones( + (num_reqs, self.num_speculative_steps + 1), + device=hidden_states.device, + dtype=torch.int64, + ) + num_logits = torch.tensor( + input_batch.cu_num_logits_np[1 : num_reqs + 1] + - input_batch.cu_num_logits_np[:num_reqs], + device=hidden_states.device, + dtype=torch.int32, + ) + accepted = torch.full_like(num_logits, num_accepted_tokens) + num_sampled = torch.minimum(accepted, num_logits) + prefill_lens = self.req_states.prefill_len.gpu[input_batch.idx_mapping] + is_chunked_prefill = input_batch.seq_lens[:num_reqs] < prefill_lens + num_sampled = torch.where(is_chunked_prefill, 0, num_sampled) + num_rejected = torch.where(is_chunked_prefill, 0, num_logits - num_sampled) + sampler_output = MRV2SamplerOutput( + sampled_token_ids=sampled_token_ids, + logprobs_tensors=None, + num_nans=None, + num_sampled=num_sampled, + ) + return sampler_output, num_sampled, num_rejected + + monkeypatch.setattr( + InprocClient, + "get_output", + get_fake_step_action_fn(original_step_action_fn), + ) + monkeypatch.setattr( + KVCacheManager, + "allocate_slots", + get_fake_allocate_slots_fn(original_allocate_slots), + ) + monkeypatch.setattr(MRV2GPUModelRunner, "execute_model", wrapped_execute_model) + monkeypatch.setattr(MRV2GPUModelRunner, "sample", fake_sample) + monkeypatch.setattr( + MambaHybridModelState, "preprocess_state", wrapped_preprocess_state + ) + monkeypatch.setattr( + MambaHybridModelState, "postprocess_state", wrapped_postprocess_state + ) + + engine = LLM( + model=MODEL, + load_format="dummy", + enforce_eager=True, + skip_tokenizer_init=True, + enable_prefix_caching=True, + block_size=BLOCK_SIZE, + mamba_cache_mode="align", + speculative_config={ + "method": "qwen3_next_mtp", + "num_speculative_tokens": num_speculative_tokens, + }, + max_num_batched_tokens=3072, + max_model_len=BLOCK_SIZE * 12, + hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS}, + seed=42, + ) + + try: + tests = get_mamba_prefix_cache_step_configs() + + global step_actions + global cur_step_action_idx + global num_accepted_tokens + for test_name, test_config in tests.items(): + num_accepted_tokens = test_config.num_accepted_tokens + cur_step_action_idx = 0 + fill_following_kv_cache_block_ids(test_config) + step_actions = test_config.step_actions + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=test_config.num_generated_tokens, + ignore_eos=True, + ) + _ = engine.generate( + [TokensPrompt(prompt_token_ids=[1] * test_config.num_prompt_tokens)], + sampling_params=sampling_params, + ) + assert cur_step_action_idx == len(test_config.step_actions), test_name + assert ( + engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache() + ) + + step_actions = [] + cur_step_action_idx = 0 + num_accepted_tokens = 1 + prompt = TokensPrompt(prompt_token_ids=[1] * (BLOCK_SIZE * 2)) + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=1, + ignore_eos=True, + ) + _ = engine.generate([prompt], sampling_params=sampling_params) + first_event_count = len(events) + _ = engine.generate([prompt], sampling_params=sampling_params) + second_events = events[first_event_count:] + prefix_hits = [ + num_computed_tokens + for num_computed_tokens in second_events + if num_computed_tokens >= BLOCK_SIZE + ] + assert prefix_hits, ( + "Expected the second identical prompt to hit prefix cache, " + f"got events={second_events!r}" + ) + assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache() + finally: + del engine + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93b2..b093b1788a9 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1997,13 +1997,6 @@ class VllmConfig: model_config = self.model_config speculative_config = self.speculative_config - if ( - model_config is not None - and model_config.has_inner_state - and self.cache_config.mamba_cache_mode == "align" - ): - unsupported.append("hybrid/mamba models with align cache mode") - if self.parallel_config.prefill_context_parallel_size > 1: unsupported.append("prefill context parallelism") @@ -2152,10 +2145,6 @@ class VllmConfig: "to schedule a multiple of block_size tokens even if they are " "in the middle of a mm input" ) - # TODO: support align mamba cache mode for model runner v2 - assert not envs.VLLM_USE_V2_MODEL_RUNNER, ( - "Model Runner V2 has not yet supported mamba_cache_mode='align'. " - ) @model_validator(mode="after") def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig": diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 6121e55dab8..eebb5ef148e 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -973,7 +973,9 @@ class DiffusionGemmaModelState(ModelState): # so the captured graph and runtime point to identical addresses. return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]} - def postprocess_state(self, idx_mapping, num_sampled) -> None: + def postprocess_state( + self, idx_mapping, num_sampled, num_computed_tokens=None + ) -> None: return None def prepare_attn( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8869e93f9fa..c9f4362a6fb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1111,7 +1111,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.req_states.total_len.gpu, ) - self.model_state.postprocess_state(idx_mapping, num_sampled) + self.model_state.postprocess_state( + idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu + ) @torch.inference_mode() def execute_model( @@ -1176,6 +1178,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Prepare all the inputs and copy to the input buffers. input_batch = self.prepare_inputs(scheduler_output, batch_desc) block_tables, slot_mappings = self.prepare_attn(input_batch) + # Mamba "align" pre-copy: migrate recurrent state across block + # boundaries before the forward. Runs only on real batches, and + # before model_state.prepare_attn gathers num_accepted_tokens so the + # boundary reset is visible to the attention metadata. + self.model_state.preprocess_state( + input_batch, + block_tables, + self.kv_cache_config, + self.req_states.num_computed_tokens.gpu, + ) if self.lora_config: # Activate LoRA adapters. diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index c80e19547c0..df86efa4a79 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -95,8 +95,23 @@ class ModelState(ABC): def apply_staged_writes(self) -> None: return None + def preprocess_state( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + kv_cache_config: KVCacheConfig, + num_computed_tokens: torch.Tensor, + ) -> None: + """Hook run on real batches before the forward pass (after block tables + are gathered). Used by mamba "align" prefix caching to pre-copy state + across block boundaries. No-op by default.""" + return None + def postprocess_state( - self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor, + num_computed_tokens: torch.Tensor | None = None, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index a0f5968361c..c6a0632c2d3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -9,15 +9,25 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.layers.mamba.mamba_utils import ( + get_conv_copy_spec, + is_conv_state_dim_first, +) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata +from vllm.v1.worker.mamba_utils import ( + MambaSpecDecodeGPUContext, + preprocess_mamba_align_fused_kernel, +) from vllm.v1.worker.utils import AttentionGroup @@ -65,9 +75,142 @@ class MambaHybridModelState(DefaultModelState): device: torch.device, ) -> None: super().__init__(vllm_config, model, encoder_cache, device) + self.cache_config = vllm_config.cache_config self.num_accepted_tokens_gpu = torch.ones( self.max_num_reqs, dtype=torch.int32, device=self.device ) + # Pre-copy "align" prefix-cache state (V2). The migration of each + # request's mamba state across block boundaries runs as a fused GPU + # kernel reusing the postprocess copy machinery, so the per-step src + # columns and the running state_idx are kept GPU-resident. + self._align_mode = self.cache_config.mamba_cache_mode == "align" + if self._align_mode: + self._mamba_state_idx_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._mamba_src_col_gpu = torch.full( + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device + ) + self._mamba_src_off_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._mamba_ctx: MambaSpecDecodeGPUContext | None = None + self._mamba_group_ids: list[int] = [] + self._mamba_spec: MambaSpec | None = None + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + if self._align_mode: + # Seed the running state block from the resumed/prefilled position. + self._mamba_state_idx_gpu[req_index] = ( + new_req_data.num_computed_tokens - 1 + ) // self.cache_config.block_size + self.num_accepted_tokens_gpu[req_index] = 1 + + def _get_mamba_group_info( + self, kv_cache_config: KVCacheConfig + ) -> tuple[list[int], MambaSpec]: + if self._mamba_spec is None: + group_ids: list[int] = [] + specs: list[MambaSpec] = [] + for i, group in enumerate(kv_cache_config.kv_cache_groups): + spec = group.kv_cache_spec + if isinstance(spec, MambaSpec): + group_ids.append(i) + specs.append(spec) + assert specs, "no mamba layers in the model" + assert all(specs[0] == s for s in specs) + self._mamba_group_ids = group_ids + self._mamba_spec = specs[0] + return self._mamba_group_ids, self._mamba_spec + + def _ensure_align_ctx( + self, + kv_cache_config: KVCacheConfig, + mamba_group_ids: list[int], + block_tables: tuple[torch.Tensor, ...], + ) -> MambaSpecDecodeGPUContext: + if self._mamba_ctx is None: + copy_funcs = self.model.get_mamba_state_copy_func() + # The fused copy kernels shift conv windows assuming the SD layout; + # the DS layout cannot express a >0 spec-decode shift as a single + # contiguous copy (mirrors get_conv_copy_spec's NotImplementedError). + if get_conv_copy_spec in copy_funcs and is_conv_state_dim_first(): + assert self.vllm_config.speculative_config is None, ( + "DS conv state layout does not support mamba align state " + "copies with speculative decoding" + ) + self._mamba_ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=kv_cache_config, + num_state_types=len(copy_funcs), + device=self.device, + make_buffer=lambda n, dtype: CpuGpuBuffer( + n, dtype=dtype, device=self.device + ), + ) + ctx = self._mamba_ctx + if not ctx.is_initialized: + forward_context = self.vllm_config.compilation_config.static_forward_context + # block_tables are batch-order slices of the persistent + # input_block_tables (stable data_ptr), so the metadata is captured + # once here and reused across steps. + ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + self.model.get_mamba_state_copy_func(), + [block_tables[gid] for gid in mamba_group_ids], + ) + return ctx + + def preprocess_state( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + kv_cache_config: KVCacheConfig, + num_computed_tokens: torch.Tensor, + ) -> None: + """Migrate each request's mamba state across block boundaries before the + forward (V1 align semantics, done on GPU). Runs on real batches only + (dummy DP/profiling runs skip preprocess_state), and before + ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset + is visible to the forward kernels. + """ + if not self._align_mode: + return + num_reqs = input_batch.num_reqs + if num_reqs == 0: + return + mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) + ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + + # The state-advance + pre-copy kernels run every step; they fast-exit per + # request when src_col < 0 or src_col == dst_col, so no copy happens on + # steps that don't cross a block boundary. (Skipping the launch entirely + # would need a V1-style async-D2H of the actual num_computed, since + # num_computed_tokens_np is an optimistic mirror under async scheduling; + # the launch cost is ~0.3% of TPOT, so the GPU fast-exit suffices.) + block = 256 + grid = (triton.cdiv(num_reqs, block),) + preprocess_mamba_align_fused_kernel[grid]( + input_batch.idx_mapping, + self._mamba_state_idx_gpu, + num_computed_tokens, + input_batch.query_start_loc, + self.num_accepted_tokens_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + num_reqs, + BLOCK_SIZE=block, + MAMBA_BLOCK_SIZE=mamba_spec.block_size, + ) + ctx.run_fused_precopy( + num_reqs, + self._mamba_state_idx_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + input_batch.idx_mapping, + ) def prepare_attn( self, @@ -145,22 +288,45 @@ class MambaHybridModelState(DefaultModelState): ) def postprocess_state( - self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor | int, + num_computed_tokens: torch.Tensor | None = None, ) -> None: # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. if not isinstance(num_sampled, int): # idx_mapping may contain -1 sentinels (filtered rows) under PP; the # kernel skips them rather than scattering with a host-side gather. - num_reqs = idx_mapping.shape[0] - if num_reqs: - _scatter_num_accepted_kernel[(num_reqs,)]( + n = idx_mapping.shape[0] + if n: + _scatter_num_accepted_kernel[(n,)]( idx_mapping, num_sampled, self.num_accepted_tokens_gpu ) - return + else: + # Fill with single value. + self.num_accepted_tokens_gpu.index_fill_( + 0, idx_mapping, max(num_sampled, 1) + ) - # Fill with single value. - self.num_accepted_tokens_gpu.index_fill_(0, idx_mapping, max(num_sampled, 1)) + # Align: save the running state to the block-aligned position when + # spec-decode acceptance leaves the sequence non-block-aligned (mirrors + # the V1 align postprocess). num_computed_tokens already holds the + # post-step advanced count. + if ( + self._align_mode + and num_computed_tokens is not None + and self._mamba_ctx is not None + ): + num_reqs = idx_mapping.shape[0] + if num_reqs: + self._mamba_ctx.run_fused_postprocess_align( + num_reqs, + self.num_accepted_tokens_gpu, + self._mamba_state_idx_gpu, + num_computed_tokens, + idx_mapping, + ) @triton.jit diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 3192b9aeaa3..ff9a75f0ef1 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -17,6 +17,7 @@ from vllm.v1.core.sched.output import ( NewRequestData, SchedulerOutput, ) +from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner @@ -177,9 +178,17 @@ def warmup_kernels( num_kv_cache_groups = len(kv_cache_groups) # Compute per-request block counts for each KV cache group. - group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups] - prefill_block_counts = [cdiv(prompt_len, bs) for bs in group_block_sizes] - decode_block_counts = [cdiv(decode_len, bs) for bs in group_block_sizes] + def _warmup_block_count(num_tokens: int, spec: Any) -> int: + num_blocks = cdiv(num_tokens, spec.block_size) + if isinstance(spec, MambaSpec) and spec.mamba_cache_mode == "align": + # Align mode reserves extra blocks beyond the token range for the + # speculative-decode running-state snapshots. + num_blocks += spec.num_speculative_blocks + return num_blocks + + kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] + prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs] + decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs] decode_block_deltas = [ d - p for d, p in zip(decode_block_counts, prefill_block_counts) ] diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 45166ef9a3a..8d8d3e62a9d 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -23,6 +23,112 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch +@triton.jit +def _copy_mamba_state_block( + state_idx, + bt_row_idx, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + # DS conv row metadata. Zero keeps the single-region copy path. + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + COPY_BLOCK_SIZE: tl.constexpr, + CONV_STATE_DIM_FIRST: tl.constexpr, +): + """Copy one (layer, state-type) mamba state block between block columns. + + Shared copy body of ``postprocess_mamba_fused_kernel`` and + ``precopy_mamba_align_fused_kernel``, mirroring the V1 copy specs + (``get_conv_copy_spec`` / ``get_temporal_copy_spec``): + - conv state (conv_width > 0): shift the window by ``token_bias`` tokens, + ``state[bt[src_col], token_bias:] -> + state[bt[dst_col], :conv_width - token_bias]`` + - temporal state: ``token_bias`` selects the accepted speculative column, + ``state[bt[src_col + token_bias]] -> state[bt[dst_col]]`` + + The caller owns the decision logic (which columns, whether to copy); this + device function only performs the byte copy for the given metadata slot. + """ + state_base_addr = tl.load(state_base_addrs_ptr + state_idx) + state_block_stride = tl.load(state_block_strides_ptr + state_idx) + state_elem_size = tl.load(state_elem_sizes_ptr + state_idx) + state_inner_size = tl.load(state_inner_sizes_ptr + state_idx) + conv_width = tl.load(state_conv_widths_ptr + state_idx) + + # Load the group index for this state, then index into the correct + # group's block table. Each mamba group has independently allocated + # physical blocks. Reinterpret as int32* since block ids are int32. + group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) + group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) + block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32)) + block_table_base = block_table_typed + bt_row_idx * block_table_stride_req + + # Widen block ids to int64 before they reach `block_id * state_block_stride` + # below: state_block_stride can exceed 2**31 bytes for large mamba caches, + # and Triton would otherwise do the multiply in int32 and wrap. + dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64) + dst_addr = state_base_addr + dest_block_id * state_block_stride + + is_conv_state = conv_width > 0 + + if CONV_STATE_DIM_FIRST and is_conv_state: + # DS conv layout: state_len is the slide axis; copy per dim row. + src_block_id = tl.load(block_table_base + src_col).to(tl.int64) + dim_rows = tl.load(state_dim_row_count_ptr + state_idx) + row_stride = tl.load(state_dim_row_stride_ptr + state_idx) + per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size + bias_bytes = token_bias.to(tl.int64) * state_elem_size + src_block_addr = state_base_addr + src_block_id * state_block_stride + offsets = tl.arange(0, COPY_BLOCK_SIZE) + for d in range(0, dim_rows): + row_src = src_block_addr + d * row_stride + bias_bytes + row_dst = dst_addr + d * row_stride + for i in range(0, per_row_bytes, COPY_BLOCK_SIZE): + mask = (i + offsets) < per_row_bytes + curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8)) + curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8)) + data = tl.load(curr_src, mask=mask) + tl.store(curr_dst, data, mask=mask) + return + + if is_conv_state: + # SD conv: copy + # state[bt[src_col], token_bias:] -> + # state[bt[dst_col], :conv_width - token_bias] + src_block_id = tl.load(block_table_base + src_col).to(tl.int64) + src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size + src_addr = state_base_addr + src_block_id * state_block_stride + src_offset + num_elems_to_copy = (conv_width - token_bias).to(tl.int64) * state_inner_size + copy_size = num_elems_to_copy * state_elem_size + else: + # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] + actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to( + tl.int64 + ) + src_addr = state_base_addr + actual_src_block_id * state_block_stride + # Use natural block data size (inner_size * elem_size), NOT + # state_block_stride which is the page stride and can exceed the + # actual data when the state tensor uses as_strided page padding. + copy_size = state_inner_size * state_elem_size + + offsets = tl.arange(0, COPY_BLOCK_SIZE) + for i in range(0, copy_size, COPY_BLOCK_SIZE): + mask = (i + offsets) < copy_size + curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8)) + curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8)) + data = tl.load(curr_src, mask=mask) + tl.store(curr_dst, data, mask=mask) + + @triton.jit def postprocess_mamba_fused_kernel( # Decision inputs (per-request) @@ -49,6 +155,10 @@ def postprocess_mamba_fused_kernel( state_dim_row_stride_ptr, # int64: bytes between rows for DS conv # Output: num_accepted_tokens update (for src==dst case) num_accepted_tokens_out_ptr, + # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The + # per-request decision arrays are in req-state-slot order; the block table + # is in batch order, so HAS_IDX_MAPPING splits the two indexings. + idx_mapping_ptr, # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation) num_reqs, # Compile-time constants (fixed after model initialization) @@ -57,35 +167,53 @@ def postprocess_mamba_fused_kernel( # COPY_BLOCK_SIZE: fixed tuning parameter for memory copy loop COPY_BLOCK_SIZE: tl.constexpr, CONV_STATE_DIM_FIRST: tl.constexpr, + # HAS_IDX_MAPPING: when True, program_id(0) is a batch index resolved to a + # req-state slot via idx_mapping_ptr (V2). When False, it is the req index. + HAS_IDX_MAPPING: tl.constexpr = False, + # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds + # the post-step new_num_computed value (V2 supplies the advanced count). + PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, ): """ Fused GPU kernel for postprocess_mamba that computes decisions AND performs mamba state copies without any CPU-GPU synchronization. Grid: (num_reqs, num_layers * num_state_types) - - program_id(0) = request index + - program_id(0) = request/batch index - program_id(1) = state_idx (flattened index into layer/state_type metadata) Note: num_layers and num_state_types are not passed as kernel parameters because the kernel indexes directly into pre-flattened metadata arrays using program_id(1). The grid dimensions encode the total state count. """ - req_idx = tl.program_id(0) + batch_idx = tl.program_id(0) state_idx = tl.program_id(1) # Bounds check - if req_idx >= num_reqs: + if batch_idx >= num_reqs: return + if HAS_IDX_MAPPING: + req_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_idx < 0: + return + else: + req_idx = batch_idx + # Compute decision logic (mirrors postprocess_mamba Python reference) num_accepted = tl.load(num_accepted_tokens_ptr + req_idx) src_block_idx = tl.load(mamba_state_idx_ptr + req_idx) - num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx) - num_computed = tl.load(num_computed_tokens_ptr + req_idx) - num_draft = tl.load(num_draft_tokens_ptr + req_idx) - num_tokens_running_state = num_computed + num_scheduled - num_draft - new_num_computed = num_tokens_running_state + num_accepted - 1 + if PRECOMPUTED_NEW_COMPUTED: + new_num_computed = tl.load(num_computed_tokens_ptr + req_idx) + num_tokens_running_state = new_num_computed - num_accepted + 1 + else: + num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx) + num_computed = tl.load(num_computed_tokens_ptr + req_idx) + num_draft = tl.load(num_draft_tokens_ptr + req_idx) + num_tokens_running_state = num_computed + num_scheduled - num_draft + new_num_computed = num_tokens_running_state + num_accepted - 1 + aligned_new_computed = (new_num_computed // block_size) * block_size needs_copy = aligned_new_computed >= num_tokens_running_state @@ -97,99 +225,158 @@ def postprocess_mamba_fused_kernel( accept_token_bias = aligned_new_computed - num_tokens_running_state dest_block_idx = aligned_new_computed // block_size - 1 - # Load state metadata for this layer/state_type - state_base_addr = tl.load(state_base_addrs_ptr + state_idx) - state_block_stride = tl.load(state_block_strides_ptr + state_idx) - state_elem_size = tl.load(state_elem_sizes_ptr + state_idx) - state_inner_size = tl.load(state_inner_sizes_ptr + state_idx) - conv_width = tl.load(state_conv_widths_ptr + state_idx) - - # Load the group index for this state, then index into the correct - # group's block table. Each mamba group has independently allocated - # physical blocks. - group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) - - # block_table_ptrs_ptr holds one pointer per group (each group owns its own - # block table). Reinterpret as int32* since block ids are int32. - group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) - block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32)) - block_table_base = block_table_typed + req_idx * block_table_stride_req - - # Widen block ids to int64 before they reach `block_id * state_block_stride` - # below: state_block_stride can exceed 2**31 bytes for large mamba caches, - # and Triton would otherwise do the multiply in int32 and wrap. - src_block_id = tl.load(block_table_base + src_block_idx).to(tl.int64) - dest_block_id = tl.load(block_table_base + dest_block_idx).to(tl.int64) - - # Compute source and destination addresses based on state type - # conv_width > 0 means this is a conv state (get_conv_copy_spec logic) - # conv_width == 0 means this is a temporal state (get_temporal_copy_spec logic) - is_conv_state = conv_width > 0 - - # Update accepted-token count before early exits. + # Update accepted-token count before early exits (per-request, so only + # state_idx == 0 writes). V2 updates in place; V1 writes the _out buffer. if src_block_idx == dest_block_idx and state_idx == 0: - tl.store(num_accepted_tokens_out_ptr + req_idx, 1) + if HAS_IDX_MAPPING: + tl.store(num_accepted_tokens_ptr + req_idx, 1) + else: + tl.store(num_accepted_tokens_out_ptr + req_idx, 1) # Skip no-op self-copy. if src_block_idx == dest_block_idx and accept_token_bias == 0: return - if CONV_STATE_DIM_FIRST and is_conv_state: - dim_rows = tl.load(state_dim_row_count_ptr + state_idx) - row_stride = tl.load(state_dim_row_stride_ptr + state_idx) - per_row_bytes = (conv_width - accept_token_bias).to(tl.int64) * state_elem_size - bias_bytes = accept_token_bias.to(tl.int64) * state_elem_size - src_block_addr = state_base_addr + src_block_id * state_block_stride - dst_block_addr = state_base_addr + dest_block_id * state_block_stride - offsets = tl.arange(0, COPY_BLOCK_SIZE) - for d in range(0, dim_rows): - row_src = src_block_addr + d * row_stride + bias_bytes - row_dst = dst_block_addr + d * row_stride - for i in range(0, per_row_bytes, COPY_BLOCK_SIZE): - mask = (i + offsets) < per_row_bytes - curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8)) - curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(curr_src, mask=mask) - tl.store(curr_dst, data, mask=mask) + bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx + _copy_mamba_state_block( + state_idx, + bt_row_idx, + src_block_idx, + dest_block_idx, + accept_token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + COPY_BLOCK_SIZE, + CONV_STATE_DIM_FIRST, + ) + + +@triton.jit +def preprocess_mamba_align_fused_kernel( + idx_mapping_ptr, + state_idx_ptr, + num_computed_tokens_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + src_col_ptr, + src_off_ptr, + num_reqs, + BLOCK_SIZE: tl.constexpr, + MAMBA_BLOCK_SIZE: tl.constexpr, +): + """Fused align preprocess: emit the pre-copy src column/offset AND advance + state_idx (with accepted-token reset) in a single launch (V2 align). + + Per batch_idx (0..num_reqs-1), resolving req slot via idx_mapping: + 1. Read pre-advance state_idx and num_accepted (last step's values). + 2. Store the pre-copy src columns for ``precopy_mamba_align_fused_kernel``: + - src_col = state_idx (the previous running block column) + - src_off = max(num_accepted - 1, 0) (the accepted-token bias) + 3. Advance state_idx to the new running block, and reset num_accepted to 1 + when a block boundary is crossed (so the migrated state, now at the + start of the new block, is read with the neutral bias). + """ + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_reqs + req_indices = tl.load(idx_mapping_ptr + offsets, mask=mask, other=0) + + state_idx = tl.load(state_idx_ptr + req_indices, mask=mask, other=-1) + num_accepted = tl.load(num_accepted_tokens_ptr + req_indices, mask=mask, other=1) + + src_off = tl.maximum(num_accepted - 1, 0) + tl.store(src_col_ptr + req_indices, state_idx, mask=mask) + tl.store(src_off_ptr + req_indices, src_off, mask=mask) + + num_computed = tl.load(num_computed_tokens_ptr + req_indices, mask=mask, other=0) + query_start = tl.load(query_start_loc_ptr + offsets, mask=mask, other=0) + query_end = tl.load(query_start_loc_ptr + offsets + 1, mask=mask, other=0) + computed_after = num_computed + query_end - query_start + new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 + tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask) + should_reset = (state_idx >= 0) & (state_idx != new_state_idx) + tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) + + +@triton.jit +def precopy_mamba_align_fused_kernel( + # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by + # the V2 fused align preprocess kernel for the current step: + mamba_state_idx_ptr, # post-advance dst block column + src_col_ptr, # pre-advance src block column (-1 = fresh) + token_bias_ptr, # accepted-token bias = num_accepted - 1 (pre-reset) + # Same flattened state-layout metadata as postprocess_mamba_fused_kernel + block_table_ptrs_ptr, + block_table_stride_req: tl.int64, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) + num_reqs, + COPY_BLOCK_SIZE: tl.constexpr, + CONV_STATE_DIM_FIRST: tl.constexpr, +): + """Pre-copy mamba "align" state across block boundaries on the V2 runner. + + Before the forward pass, copy each request's last SSM/conv state from its + previous block column into the new window block column, so the kernels read + the initial state from the write-side block as usual (V1 align semantics). + Same per-(layer, state) copy semantics as ``postprocess_mamba_fused_kernel`` + (shared ``_copy_mamba_state_block`` body, i.e. the V1 ``preprocess_mamba`` + copy specs), but driven by the GPU-resident src columns so it needs no + CPU-GPU sync (async-scheduling safe). + + Grid: (num_reqs, num_layers * num_state_types); block tables are indexed by + batch row, per-request state by req_idx via idx_mapping (V2 layout). + """ + batch_idx = tl.program_id(0) + state_idx = tl.program_id(1) + if batch_idx >= num_reqs: + return + req_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_idx < 0: return - if is_conv_state: - # SD conv: copy - # state[block_table[req_idx, src_block_idx], accept_token_bias:] - # to - # state[block_table[req_idx, dest_block_idx], :conv_width - accept_token_bias] - src_offset = accept_token_bias.to(tl.int64) * state_inner_size * state_elem_size - src_addr = state_base_addr + src_block_id * state_block_stride + src_offset - dst_addr = state_base_addr + dest_block_id * state_block_stride - # Number of elements to copy: - # (conv_width - accept_token_bias) * inner_size - num_elems_to_copy = (conv_width - accept_token_bias).to( - tl.int64 - ) * state_inner_size - copy_size = num_elems_to_copy * state_elem_size - else: - # Temporal state: copy - # state[block_table[req_idx, src_block_idx + accept_token_bias]] - # to - # state[block_table[req_idx, dest_block_idx]] - actual_src_block_idx = src_block_idx + accept_token_bias - actual_src_block_id = tl.load(block_table_base + actual_src_block_idx).to( - tl.int64 - ) - src_addr = state_base_addr + actual_src_block_id * state_block_stride - dst_addr = state_base_addr + dest_block_id * state_block_stride - # Use natural block data size (inner_size * elem_size), NOT - # state_block_stride which is the page stride and can exceed the - # actual data when the state tensor uses as_strided page padding. - copy_size = state_inner_size * state_elem_size + src_col = tl.load(src_col_ptr + req_idx) + dst_col = tl.load(mamba_state_idx_ptr + req_idx) + # Fresh state, or still writing the same block: kernels locate the initial + # state in-block via num_accepted (preserved when no boundary is crossed), + # so there is nothing to copy. + if src_col < 0 or src_col == dst_col: + return - offsets = tl.arange(0, COPY_BLOCK_SIZE) - for i in range(0, copy_size, COPY_BLOCK_SIZE): - mask = (i + offsets) < copy_size - curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8)) - curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(curr_src, mask=mask) - tl.store(curr_dst, data, mask=mask) + token_bias = tl.load(token_bias_ptr + req_idx) + _copy_mamba_state_block( + state_idx, + batch_idx, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + COPY_BLOCK_SIZE, + CONV_STATE_DIM_FIRST, + ) @triton.jit @@ -559,12 +746,101 @@ class MambaSpecDecodeGPUContext: self.state_dim_row_count, self.state_dim_row_stride, self.num_accepted_tokens_out, + None, # idx_mapping: V1 decision arrays are already in req order num_reqs, block_size=self.block_size, COPY_BLOCK_SIZE=1024, CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), ) + def run_fused_precopy( + self, + num_reqs: int, + state_idx_gpu: torch.Tensor, + src_col_gpu: torch.Tensor, + token_bias_gpu: torch.Tensor, + idx_mapping: torch.Tensor, + ) -> None: + """Pre-copy each request's previous running block into its new window + block before the forward pass (V2 align boundary migration). + + Args: + num_reqs: Number of active requests (batch order). + state_idx_gpu: [max_reqs] post-advance dst block column per req slot. + src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh). + token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1). + idx_mapping: [num_reqs] batch_idx -> req_state_idx (-1 to skip). + """ + if num_reqs == 0 or not self.is_initialized: + return + total_states = self.num_layers * self.num_state_types + grid = (num_reqs, total_states) + precopy_mamba_align_fused_kernel[grid]( + state_idx_gpu, + src_col_gpu, + token_bias_gpu, + self.block_table_ptrs, + self.block_table_stride_req, + self.state_base_addrs, + self.state_block_strides, + self.state_elem_sizes, + self.state_inner_sizes, + self.state_conv_widths, + self.state_group_indices, + self.state_dim_row_count, + self.state_dim_row_stride, + idx_mapping, + num_reqs, + COPY_BLOCK_SIZE=1024, + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), + ) + + def run_fused_postprocess_align( + self, + num_reqs: int, + num_accepted_tokens_gpu: torch.Tensor, + state_idx_gpu: torch.Tensor, + new_num_computed_tokens_gpu: torch.Tensor, + idx_mapping: torch.Tensor, + ) -> None: + """V2 align postprocess: save the running state to the block-aligned + position after spec-decode acceptance leaves the sequence non-aligned. + + ``num_accepted_tokens_gpu`` is updated in place (reset to 1 when the + accepted position stays in the running block); ``new_num_computed_tokens`` + already holds the post-step computed count (PRECOMPUTED_NEW_COMPUTED). + ``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING). + """ + if num_reqs == 0 or not self.is_initialized: + return + total_states = self.num_layers * self.num_state_types + grid = (num_reqs, total_states) + postprocess_mamba_fused_kernel[grid]( + num_accepted_tokens_gpu, + state_idx_gpu, + None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED + new_num_computed_tokens_gpu, + None, # num_draft: unused under PRECOMPUTED_NEW_COMPUTED + self.block_table_ptrs, + self.block_table_stride_req, + self.state_base_addrs, + self.state_block_strides, + self.state_elem_sizes, + self.state_inner_sizes, + self.state_conv_widths, + self.state_group_indices, + self.state_dim_row_count, + self.state_dim_row_stride, + None, # num_accepted_out: V2 updates num_accepted in place + idx_mapping, + num_reqs, + block_size=self.block_size, + COPY_BLOCK_SIZE=1024, + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), + HAS_IDX_MAPPING=True, + PRECOMPUTED_NEW_COMPUTED=True, + ) + @dataclasses.dataclass class MambaBuffers: From 5316638a5eb98d764a5618c20a1558ffc24d3bc9 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 29 Jun 2026 22:20:33 +0100 Subject: [PATCH 132/138] Fix transient dependency issues caused by `requirements/common.txt` (#47015) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docker/Dockerfile.cpu | 1 + requirements/test/cuda.in | 6 +- requirements/test/cuda.txt | 306 ++++++++++++++++- requirements/test/rocm.in | 4 - requirements/test/rocm.txt | 4 - requirements/test/xpu.in | 2 + requirements/test/xpu.txt | 320 +++++++++++++++++- .../test_structural_tag_registry.py | 67 ++-- 8 files changed, 638 insertions(+), 72 deletions(-) diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 61bad68b442..adb94b5a927 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -194,6 +194,7 @@ FROM base AS vllm-test-deps WORKDIR /vllm-workspace # Copy test requirements +COPY requirements/common.txt requirements/common.txt COPY requirements/test/cuda.in requirements/test/cpu.in RUN \ diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 12a40716392..9a6e46712cb 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -1,3 +1,5 @@ +-r ../common.txt + # testing pytest tensorizer==2.10.1 @@ -13,7 +15,6 @@ albumentations # required for Nemotron Parse in test_common.py av # required for audio_in_video tests backoff # required for phi4mm test blobfile # required for kimi-vl test -einops # required for MPT, qwen-vl httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test @@ -34,7 +35,6 @@ matplotlib # required for qwen-vl test mistral_common[image,audio] >= 1.11.5 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py -opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test @@ -55,11 +55,9 @@ grpcio-reflection==1.78.0 arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test numba == 0.65.0 # Required for N-gram speculative decoding -numpy runai-model-streamer[s3,gcs,azure]==0.15.7 fastsafetensors>=0.3.2 instanttensor>=0.1.5; platform_machine == "x86_64" -pydantic>=2.12 # 2.11 leads to error on python 3.13 decord==0.6.0; platform_machine == "x86_64" # terratorch is temporarily disabled while PyPI has the `lightning` package # in `quarantined` status (every published terratorch version transitively diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index f504c69c48f..e8d600ba632 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -9,6 +9,7 @@ aiohappyeyeballs==2.6.1 aiohttp==3.13.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # aiohttp-cors # datasets # fsspec @@ -24,17 +25,34 @@ albumentations==1.4.6 alembic==1.16.4 # via optuna annotated-doc==0.0.4 - # via fastapi + # via + # fastapi + # typer annotated-types==0.7.0 # via pydantic -anyio==4.6.2.post1 +anthropic==0.112.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +anyio==4.14.1 + # via + # anthropic # httpx + # mcp + # openai + # sse-starlette # starlette + # watchfiles +apache-tvm-ffi==0.1.9 + # via + # -c requirements/cuda.txt + # xgrammar arctic-inference==0.1.1 # via -r requirements/test/cuda.in argcomplete==3.5.1 # via datamodel-code-generator +astor==0.8.1 + # via depyf attrs==24.2.0 # via # aiohttp @@ -59,6 +77,8 @@ bitsandbytes==0.49.2 # via -r requirements/test/cuda.in black==24.10.0 # via datamodel-code-generator +blake3==1.0.9 + # via -r requirements/test/../common.txt blobfile==3.0.0 # via -r requirements/test/cuda.in bm25s==0.2.13 @@ -76,12 +96,17 @@ bounded-pool-executor==0.0.3 buildkite-test-collector==0.1.9 # via -r requirements/test/cuda.in cachetools==5.5.2 - # via google-auth + # via + # -r requirements/test/../common.txt + # google-auth +cbor2==6.1.2 + # via -r requirements/test/../common.txt certifi==2024.8.30 # via # httpcore # httpx # requests + # sentry-sdk cffi==2.0.0 # via # cryptography @@ -98,9 +123,11 @@ click==8.1.7 # jiwer # nltk # ray + # rich-toolkit # schemathesis - # typer # uvicorn +cloudpickle==3.1.2 + # via -r requirements/test/../common.txt cohere-melody==0.9.0 # via -r requirements/test/cuda.in colorama==0.4.6 @@ -111,6 +138,10 @@ colorful==0.5.6 # via ray colorlog==6.10.1 # via optuna +compressed-tensors==0.17.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt contourpy==1.3.0 # via matplotlib coverage==7.10.6 @@ -149,30 +180,49 @@ decorator==5.1.1 # via librosa decord==0.6.0 # via -r requirements/test/cuda.in +depyf==0.20.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +detect-installer==0.1.0 + # via fastapi-cloud-cli dill==0.3.8 # via # datasets + # depyf # evaluate # lm-eval # multiprocess +diskcache==5.6.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt distlib==0.3.9 # via virtualenv +distro==1.9.0 + # via + # anthropic + # openai dnspython==2.7.0 # via email-validator docker==7.1.0 # via gpt-oss docopt==0.6.2 # via num2words +docstring-parser==0.18.0 + # via anthropic einops==0.8.1 # via - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # encodec # vector-quantize-pytorch # vocos einx==0.3.0 # via vector-quantize-pytorch email-validator==2.2.0 - # via pydantic + # via + # fastapi + # pydantic encodec==0.1.1 # via vocos et-xmlfile==2.0.0 @@ -182,7 +232,17 @@ evaluate==0.4.3 fastapi==0.136.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss + # model-hosting-container-standards +fastapi-cli==0.0.27 + # via fastapi +fastapi-cloud-cli==0.21.0 + # via fastapi-cli +fastar==0.11.0 + # via + # fastapi + # fastapi-cloud-cli fastparquet==2024.11.0 # via genai-perf fastrlock==0.8.2 @@ -194,6 +254,7 @@ fastsafetensors==0.3.2 filelock==3.16.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # blobfile # datasets # huggingface-hub @@ -243,7 +304,10 @@ google-crc32c==1.7.1 google-resumable-media==2.7.2 # via google-cloud-storage googleapis-common-protos==1.70.0 - # via google-api-core + # via + # google-api-core + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http gpt-oss==0.0.8 # via -r requirements/test/cuda.in graphql-core==3.2.6 @@ -254,6 +318,7 @@ grpcio==1.78.0 # via # -r requirements/test/cuda.in # grpcio-reflection + # opentelemetry-exporter-otlp-proto-grpc # ray grpcio-reflection==1.78.0 # via -r requirements/test/cuda.in @@ -275,12 +340,22 @@ html2text==2025.4.15 # via gpt-oss httpcore==1.0.6 # via httpx +httptools==0.8.0 + # via uvicorn httpx==0.27.2 # via # -r requirements/test/cuda.in + # anthropic + # fastapi + # fastapi-cloud-cli # huggingface-hub + # mcp + # model-hosting-container-standards + # openai # perceptron # schemathesis +httpx-sse==0.4.3 + # via mcp huggingface-hub==1.10.2 # via # accelerate @@ -314,6 +389,8 @@ idna==3.10 # httpx # requests # yarl +ijson==3.5.0 + # via -r requirements/test/../common.txt imagehash==4.3.2 # via -r requirements/test/cuda.in imageio==2.37.0 @@ -326,6 +403,8 @@ iniconfig==2.0.0 # via pytest instanttensor==0.1.5 # via -r requirements/test/cuda.in +interegular==0.3.3 + # via lm-format-enforcer isodate==0.7.2 # via azure-storage-blob isort==5.13.2 @@ -333,15 +412,21 @@ isort==5.13.2 jinja2==3.1.6 # via # datamodel-code-generator + # fastapi # genai-perf # lm-eval # torch +jiter==0.15.0 + # via + # anthropic + # openai jiwer==3.0.5 # via -r requirements/test/cuda.in jmespath==1.0.1 # via # boto3 # botocore + # model-hosting-container-standards joblib==1.4.2 # via # librosa @@ -350,7 +435,9 @@ joblib==1.4.2 jsonschema==4.23.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # hypothesis-jsonschema + # mcp # mistral-common # ray jsonschema-rs==0.46.5 @@ -365,6 +452,10 @@ kaleido==0.2.1 # via genai-perf kiwisolver==1.4.7 # via matplotlib +lark==1.2.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt lazy-loader==0.4 # via # librosa @@ -373,10 +464,20 @@ libnacl==2.1.0 # via tensorizer librosa==0.10.2.post1 # via -r requirements/test/cuda.in +llguidance==1.7.6 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt llvmlite==0.47.0 # via numba lm-eval==0.4.12 # via -r requirements/test/cuda.in +lm-format-enforcer==0.11.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +loguru==0.7.3 + # via compressed-tensors lxml==5.3.0 # via # blobfile @@ -398,12 +499,19 @@ mbstrdecoder==1.1.3 # dataproperty # pytablewriter # typepy +mcp==1.28.1 + # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py mistral-common==1.11.5 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/cuda.in +model-hosting-container-standards==0.1.16 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt more-itertools==10.5.0 # via lm-eval mpmath==1.3.0 @@ -418,6 +526,8 @@ msgpack==1.1.0 # via # librosa # ray +msgspec==0.21.1 + # via -r requirements/test/../common.txt mteb==2.8.3 # via -r requirements/test/cuda.in multidict==6.1.0 @@ -434,6 +544,8 @@ networkx==3.2.1 # via # scikit-image # torch +ninja==1.13.0 + # via -r requirements/test/../common.txt nltk==3.9.1 # via rouge-score num2words==0.5.14 @@ -445,7 +557,7 @@ numba==0.65.0 # librosa numpy==2.2.6 # via - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # accelerate # albumentations # bitsandbytes @@ -489,6 +601,7 @@ numpy==2.2.6 # transformers # tritonclient # vocos + # xgrammar nvidia-cublas==13.1.0.3 # via # cuda-toolkit @@ -530,9 +643,14 @@ nvidia-nvtx==13.0.85 # via cuda-toolkit open-clip-torch==2.32.0 # via -r requirements/test/cuda.in +openai==2.44.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt openai-harmony==0.0.4 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss opencensus==0.11.4 # via ray @@ -541,7 +659,7 @@ opencensus-context==0.1.3 opencv-python-headless==4.13.0.90 # via # -c requirements/common.txt - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # albumentations # mistral-common openpyxl==3.1.5 @@ -549,24 +667,54 @@ openpyxl==3.1.5 opentelemetry-api==1.35.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http # opentelemetry-exporter-prometheus # opentelemetry-sdk # opentelemetry-semantic-conventions +opentelemetry-exporter-otlp==1.35.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +opentelemetry-exporter-otlp-proto-common==1.35.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-exporter-otlp-proto-grpc==1.35.0 + # via opentelemetry-exporter-otlp +opentelemetry-exporter-otlp-proto-http==1.35.0 + # via opentelemetry-exporter-otlp opentelemetry-exporter-prometheus==0.56b0 # via ray opentelemetry-proto==1.35.0 - # via ray + # via + # opentelemetry-exporter-otlp-proto-common + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # ray opentelemetry-sdk==1.35.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http # opentelemetry-exporter-prometheus # ray opentelemetry-semantic-conventions==0.56b0 # via opentelemetry-sdk +opentelemetry-semantic-conventions-ai==0.4.13 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt optuna==3.6.1 # via genai-perf orjson==3.11.5 # via genai-perf +outlines-core==0.2.14 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt packaging==24.2 # via # accelerate @@ -578,6 +726,7 @@ packaging==24.2 # fastparquet # huggingface-hub # lazy-loader + # lm-format-enforcer # matplotlib # optuna # peft @@ -597,6 +746,8 @@ pandas==2.2.3 # fastparquet # genai-perf # statsmodels +partial-json-parser==0.2.1.1.post7 + # via -r requirements/test/../common.txt pathspec==0.12.1 # via black pathvalidate==3.2.1 @@ -611,6 +762,7 @@ perf-analyzer==0.1.0 # via genai-perf pillow==10.4.0 # via + # -r requirements/test/../common.txt # genai-perf # imagehash # imageio @@ -644,8 +796,14 @@ pqdm==0.2.0 prometheus-client==0.22.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # opentelemetry-exporter-prometheus + # prometheus-fastapi-instrumentator # ray +prometheus-fastapi-instrumentator==8.0.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt propcache==0.2.0 # via # aiohttp @@ -655,6 +813,7 @@ proto-plus==1.26.1 protobuf==6.33.6 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # google-api-core # googleapis-common-protos # grpcio-reflection @@ -664,11 +823,14 @@ protobuf==6.33.6 # tensorizer psutil==6.1.0 # via + # -r requirements/test/../common.txt # accelerate # peft # tensorizer py==1.11.0 # via pytest-forked +py-cpuinfo==9.0.0 + # via -r requirements/test/../common.txt py-spy==0.4.0 # via ray pyarrow==23.0.0 @@ -681,6 +843,8 @@ pyasn1==0.6.1 # rsa pyasn1-modules==0.4.2 # via google-auth +pybase64==1.4.3 + # via -r requirements/test/../common.txt pycountry==24.6.1 # via pydantic-extra-types pycparser==2.22 @@ -690,26 +854,43 @@ pycryptodomex==3.22.0 pydantic==2.12.0 # via # -c requirements/common.txt - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # albumentations + # anthropic + # compressed-tensors # datamodel-code-generator # fastapi + # fastapi-cloud-cli # gpt-oss + # lm-format-enforcer + # mcp # mistral-common + # model-hosting-container-standards # mteb + # openai # openai-harmony # pydantic-extra-types + # pydantic-settings # ray + # xgrammar pydantic-core==2.41.1 # via pydantic pydantic-extra-types==2.10.5 - # via mistral-common + # via + # fastapi + # mistral-common +pydantic-settings==2.14.2 + # via + # fastapi + # mcp pygments==2.18.0 # via # pytest # rich pyjwt==2.11.0 - # via msal + # via + # mcp + # msal pyparsing==3.2.0 # via matplotlib pyrate-limiter==4.4.0 @@ -751,6 +932,16 @@ python-dateutil==2.9.0.post0 # matplotlib # pandas # typepy +python-dotenv==1.2.2 + # via + # pydantic-settings + # uvicorn +python-json-logger==4.1.0 + # via -r requirements/test/../common.txt +python-multipart==0.0.32 + # via + # fastapi + # mcp python-rapidjson==1.20 # via tritonclient pytrec-eval-terrier==0.5.7 @@ -763,12 +954,14 @@ pywavelets==1.9.0 # via imagehash pyyaml==6.0.2 # via + # -r requirements/test/../common.txt # accelerate # albumentations # datamodel-code-generator # datasets # genai-perf # huggingface-hub + # lm-format-enforcer # optuna # peft # ray @@ -776,7 +969,12 @@ pyyaml==6.0.2 # schemathesis # timm # transformers + # uvicorn # vocos +pyzmq==27.1.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt rapidfuzz==3.12.1 # via jiwer ray==2.48.0 @@ -789,6 +987,7 @@ referencing==0.35.1 # jsonschema-specifications regex==2026.2.28 # via + # -r requirements/test/../common.txt # nltk # open-clip-torch # sacrebleu @@ -797,6 +996,7 @@ regex==2026.2.28 requests==2.32.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # azure-core # buildkite-test-collector # datasets @@ -809,6 +1009,7 @@ requests==2.32.3 # mistral-common # msal # mteb + # opentelemetry-exporter-otlp-proto-http # pooch # ray # responses @@ -822,8 +1023,15 @@ rich==13.9.4 # genai-perf # mteb # perceptron + # rich-toolkit # schemathesis # typer +rich-toolkit==0.20.1 + # via + # fastapi-cli + # fastapi-cloud-cli +rignore==0.7.6 + # via fastapi-cloud-cli rouge-score==0.1.2 # via lm-eval rpds-py==0.20.1 @@ -847,6 +1055,7 @@ sacrebleu==2.4.3 safetensors==0.7.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # accelerate # open-clip-torch # peft @@ -882,9 +1091,17 @@ sentence-transformers==5.2.0 # via # -r requirements/test/cuda.in # mteb +sentencepiece==0.2.1 + # via -r requirements/test/../common.txt +sentry-sdk==2.63.0 + # via fastapi-cloud-cli +setproctitle==1.3.7 + # via -r requirements/test/../common.txt setuptools==77.0.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # model-hosting-container-standards # pytablewriter # torch shellingham==1.5.4 @@ -894,6 +1111,7 @@ shellingham==1.5.4 six==1.16.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # junit-xml # opencensus # python-dateutil @@ -902,8 +1120,9 @@ smart-open==7.1.0 # via ray sniffio==1.3.1 # via - # anyio + # anthropic # httpx + # openai sortedcontainers==2.4.0 # via hypothesis soundfile==0.12.1 @@ -922,10 +1141,17 @@ sqlalchemy==2.0.41 # optuna sqlitedict==2.1.0 # via lm-eval +sse-starlette==3.4.5 + # via mcp starlette==1.3.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # fastapi + # mcp + # model-hosting-container-standards + # prometheus-fastapi-instrumentator + # sse-starlette # starlette-testclient starlette-testclient==0.4.1 # via schemathesis @@ -933,6 +1159,8 @@ statsmodels==0.14.4 # via genai-perf structlog==25.4.0 # via gpt-oss +supervisor==4.3.0 + # via model-hosting-container-standards sympy==1.13.3 # via # einx @@ -962,6 +1190,7 @@ tifffile==2025.3.30 tiktoken==0.12.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss # lm-eval # mistral-common @@ -973,6 +1202,7 @@ timm==1.0.17 tokenizers==0.22.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/cuda.in # transformers torch==2.11.0+cu130 @@ -981,6 +1211,7 @@ torch==2.11.0+cu130 # -r requirements/test/cuda.in # accelerate # bitsandbytes + # compressed-tensors # encodec # instanttensor # mteb @@ -994,6 +1225,7 @@ torch==2.11.0+cu130 # torchvision # vector-quantize-pytorch # vocos + # xgrammar torchaudio==2.11.0+cu130 # via # -c requirements/cuda.txt @@ -1009,6 +1241,7 @@ torchvision==0.26.0+cu130 # timm tqdm==4.67.3 # via + # -r requirements/test/../common.txt # datasets # evaluate # huggingface-hub @@ -1016,6 +1249,7 @@ tqdm==4.67.3 # mteb # nltk # open-clip-torch + # openai # optuna # peft # pqdm @@ -1025,15 +1259,20 @@ tqdm==4.67.3 transformers==5.5.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/cuda.in + # compressed-tensors # genai-perf # peft # sentence-transformers # transformers-stream-generator + # xgrammar transformers-stream-generator==0.0.5 # via -r requirements/test/cuda.in triton==3.6.0 - # via torch + # via + # torch + # xgrammar tritonclient==2.64.0 # via -r requirements/test/cuda.in typepy==1.3.2 @@ -1041,8 +1280,10 @@ typepy==1.3.2 # dataproperty # pytablewriter # tabledata -typer==0.15.2 +typer==0.26.8 # via + # fastapi-cli + # fastapi-cloud-cli # fastsafetensors # huggingface-hub # perceptron @@ -1050,9 +1291,13 @@ typer==0.15.2 typing-extensions==4.15.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # aiosignal # albumentations # alembic + # anthropic + # anyio + # apache-tvm-ffi # azure-core # azure-identity # azure-storage-blob @@ -1062,9 +1307,13 @@ typing-extensions==4.15.0 # huggingface-hub # librosa # lm-eval + # mcp # mistral-common # mteb + # openai # opentelemetry-api + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http # opentelemetry-sdk # opentelemetry-semantic-conventions # pqdm @@ -1072,17 +1321,20 @@ typing-extensions==4.15.0 # pydantic-core # pydantic-extra-types # pytest-asyncio + # rich-toolkit # schemathesis # sentence-transformers # sqlalchemy # starlette # torch - # typer # typing-inspection + # xgrammar typing-inspection==0.4.2 # via # fastapi + # mcp # pydantic + # pydantic-settings tzdata==2024.2 # via pandas urllib3==2.2.3 @@ -1092,23 +1344,41 @@ urllib3==2.2.3 # docker # requests # responses + # sentry-sdk # tritonclient uvicorn==0.35.0 - # via gpt-oss + # via + # fastapi + # fastapi-cli + # fastapi-cloud-cli + # gpt-oss + # mcp +uvloop==0.22.1 + # via uvicorn vector-quantize-pytorch==1.21.2 # via -r requirements/test/cuda.in virtualenv==20.31.2 # via ray vocos==0.1.0 # via -r requirements/test/cuda.in +watchfiles==1.2.0 + # via + # -r requirements/test/../common.txt + # uvicorn wcwidth==0.2.13 # via ftfy +websockets==16.0 + # via uvicorn werkzeug==3.1.3 # via schemathesis word2number==1.1 # via lm-eval wrapt==1.17.2 # via smart-open +xgrammar==0.2.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt xxhash==3.5.0 # via # datasets diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 6a38f384f11..dc7f03c64f6 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -15,7 +15,6 @@ albumentations # required for Nemotron Parse in test_common.py av # required for audio_in_video tests backoff # required for phi4mm test blobfile # required for kimi-vl test -einops # required for MPT, qwen-vl httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test @@ -33,7 +32,6 @@ matplotlib # required for qwen-vl test mistral_common[image,audio]>=1.11.5 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py -opencv-python-headless>=4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test @@ -54,11 +52,9 @@ grpcio-reflection==1.78.0 arctic-inference==0.1.1 # Required for suffix decoding test numba==0.65.0 # Required for N-gram speculative decoding -numpy runai-model-streamer[s3,gcs,azure]==0.15.7 fastsafetensors>=0.3.2 instanttensor>=0.1.5 -pydantic>=2.12 # 2.11 leads to error on python 3.13 decord==0.6.0 # Prithvi tests diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 726aad9a672..b191705e0da 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -205,7 +205,6 @@ docstring-parser==0.17.0 einops==0.8.2 # via # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # encodec # vector-quantize-pytorch # vocos @@ -561,7 +560,6 @@ numba==0.65.0 numpy==2.2.6 # via # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # accelerate # albumentations # bitsandbytes @@ -630,7 +628,6 @@ opencv-python-headless==4.13.0.92 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # albumentations # mistral-common openpyxl==3.1.5 @@ -834,7 +831,6 @@ pydantic==2.12.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # albumentations # anthropic # compressed-tensors diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 161e2c6871f..e2f299282ed 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -1,3 +1,5 @@ +-r ../common.txt + # --- Test Infrastructure --- tblib pytest diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 2b938e3b583..16169b99863 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -11,6 +11,7 @@ aiohappyeyeballs==2.6.1 aiohttp==3.13.4 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # fsspec # gpt-oss # lm-eval @@ -24,12 +25,25 @@ annotated-doc==0.0.4 # typer annotated-types==0.7.0 # via pydantic +anthropic==0.112.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt anyio==4.13.0 # via + # anthropic # httpx + # mcp + # openai + # sse-starlette # starlette + # watchfiles +apache-tvm-ffi==0.1.12 + # via xgrammar arctic-inference==0.1.1 # via -r requirements/test/xpu.in +astor==0.8.1 + # via depyf attrs==26.1.0 # via # aiohttp @@ -39,6 +53,8 @@ audioread==3.0.1 # via # -r requirements/test/xpu.in # librosa +blake3==1.0.9 + # via -r requirements/test/../common.txt blobfile==3.0.0 # via -r requirements/test/xpu.in bm25s==0.2.13 @@ -47,13 +63,20 @@ bm25s==0.2.13 # mteb bounded-pool-executor==0.0.3 # via pqdm +cachetools==7.1.4 + # via -r requirements/test/../common.txt +cbor2==6.1.2 + # via -r requirements/test/../common.txt certifi==2026.2.25 # via # httpcore # httpx # requests + # sentry-sdk cffi==2.0.0 - # via soundfile + # via + # cryptography + # soundfile chardet==5.2.0 # via mbstrdecoder charset-normalizer==3.4.6 @@ -64,13 +87,22 @@ click==8.3.1 # via # jiwer # nltk + # rich-toolkit # schemathesis # typer # uvicorn +cloudpickle==3.1.2 + # via -r requirements/test/../common.txt colorama==0.4.6 # via sacrebleu +compressed-tensors==0.17.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt coverage==7.13.5 # via pytest-cov +cryptography==49.0.0 + # via pyjwt dataproperty==1.1.0 # via # pytablewriter @@ -82,16 +114,35 @@ datasets==4.8.4 # mteb decorator==5.2.1 # via librosa +depyf==0.20.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +detect-installer==0.1.0 + # via fastapi-cloud-cli dill==0.4.1 # via # datasets + # depyf # evaluate # lm-eval # multiprocess +diskcache==5.6.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +distro==1.9.0 + # via + # anthropic + # openai +dnspython==2.8.0 + # via email-validator docker==7.1.0 # via gpt-oss docopt==0.6.2 # via num2words +docstring-parser==0.18.0 + # via anthropic dpcpp-cpp-rt==2025.3.2 # via # onemkl-sycl-blas @@ -100,15 +151,30 @@ dpcpp-cpp-rt==2025.3.2 # onemkl-sycl-rng # onemkl-sycl-sparse # torch +einops==0.8.2 + # via -r requirements/test/../common.txt +email-validator==2.3.0 + # via + # fastapi + # pydantic evaluate==0.4.6 # via lm-eval fastapi==0.135.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss + # model-hosting-container-standards +fastapi-cli==0.0.27 + # via fastapi +fastapi-cloud-cli==0.21.0 + # via fastapi-cli +fastar==0.11.0 + # via fastapi-cloud-cli filelock==3.25.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # blobfile # datasets # huggingface-hub @@ -124,10 +190,16 @@ fsspec==2026.2.0 # evaluate # huggingface-hub # torch +googleapis-common-protos==1.75.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http gpt-oss==0.0.8 # via -r requirements/test/xpu.in graphql-core==3.2.8 # via hypothesis-graphql +grpcio==1.81.1 + # via opentelemetry-exporter-otlp-proto-grpc h11==0.16.0 # via # httpcore @@ -140,11 +212,21 @@ html2text==2025.4.15 # via gpt-oss httpcore==1.0.9 # via httpx +httptools==0.8.0 + # via uvicorn httpx==0.28.1 # via + # anthropic # datasets + # fastapi + # fastapi-cloud-cli # huggingface-hub + # mcp + # model-hosting-container-standards + # openai # schemathesis +httpx-sse==0.4.3 + # via mcp huggingface-hub==1.10.2 # via # accelerate @@ -166,9 +248,12 @@ hypothesis-jsonschema==0.23.1 idna==3.11 # via # anyio + # email-validator # httpx # requests # yarl +ijson==3.5.0 + # via -r requirements/test/../common.txt imageio==2.37.3 # via scikit-image impi-rt==2021.17.2 @@ -212,13 +297,22 @@ intel-sycl-rt==2025.3.2 # dpcpp-cpp-rt # oneccl # torch +interegular==0.3.3 + # via lm-format-enforcer jinja2==3.1.6 # via # -c requirements/xpu.txt + # fastapi # lm-eval # torch +jiter==0.15.0 + # via + # anthropic + # openai jiwer==4.0.0 # via -r requirements/test/xpu.in +jmespath==1.1.0 + # via model-hosting-container-standards joblib==1.5.3 # via # librosa @@ -227,7 +321,9 @@ joblib==1.5.3 jsonschema==4.26.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # hypothesis-jsonschema + # mcp # mistral-common # schemathesis jsonschema-rs==0.45.0 @@ -236,16 +332,30 @@ jsonschema-specifications==2025.9.1 # via jsonschema junit-xml==1.9 # via schemathesis +lark==1.2.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt lazy-loader==0.5 # via # librosa # scikit-image librosa==0.10.2.post1 # via -r requirements/test/xpu.in +llguidance==1.7.6 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt llvmlite==0.47.0 # via numba lm-eval==0.4.12 # via -r requirements/test/xpu.in +lm-format-enforcer==0.11.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +loguru==0.7.3 + # via compressed-tensors lxml==6.0.2 # via # blobfile @@ -262,11 +372,14 @@ mbstrdecoder==1.1.4 # dataproperty # pytablewriter # typepy +mcp==1.28.1 + # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py mistral-common==1.11.5 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/xpu.in mkl==2025.3.1 # via @@ -276,6 +389,10 @@ mkl==2025.3.1 # onemkl-sycl-rng # onemkl-sycl-sparse # torch +model-hosting-container-standards==0.1.16 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt modelscope==1.35.3 # via -r requirements/test/xpu.in more-itertools==10.8.0 @@ -284,6 +401,8 @@ mpmath==1.3.0 # via sympy msgpack==1.1.2 # via librosa +msgspec==0.21.1 + # via -r requirements/test/../common.txt mteb==2.12.7 # via -r requirements/test/xpu.in multidict==6.7.1 @@ -298,6 +417,8 @@ networkx==3.6.1 # via # scikit-image # torch +ninja==1.13.0 + # via -r requirements/test/../common.txt nltk==3.9.4 # via rouge-score num2words==0.5.14 @@ -308,6 +429,7 @@ numba==0.65.0 # librosa numpy==2.2.6 # via + # -r requirements/test/../common.txt # accelerate # albumentations # bm25s @@ -333,6 +455,7 @@ numpy==2.2.6 # tifffile # torchvision # transformers + # xgrammar oneccl==2021.17.2 # via # oneccl-devel @@ -356,15 +479,65 @@ onemkl-sycl-rng==2025.3.1 # via torch onemkl-sycl-sparse==2025.3.1 # via torch +openai==2.44.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt openai-harmony==0.0.8 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss opencv-python-headless==4.13.0.92 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # albumentations # mistral-common +opentelemetry-api==1.43.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions +opentelemetry-exporter-otlp==1.43.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +opentelemetry-exporter-otlp-proto-common==1.43.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-exporter-otlp-proto-grpc==1.43.0 + # via opentelemetry-exporter-otlp +opentelemetry-exporter-otlp-proto-http==1.43.0 + # via opentelemetry-exporter-otlp +opentelemetry-proto==1.43.0 + # via + # opentelemetry-exporter-otlp-proto-common + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-sdk==1.43.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-semantic-conventions-ai +opentelemetry-semantic-conventions==0.64b0 + # via + # opentelemetry-sdk + # opentelemetry-semantic-conventions-ai +opentelemetry-semantic-conventions-ai==0.5.1 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +outlines-core==0.2.14 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt packaging==26.0 # via # -c requirements/xpu.txt @@ -373,6 +546,7 @@ packaging==26.0 # evaluate # huggingface-hub # lazy-loader + # lm-format-enforcer # modelscope # pooch # pytest @@ -384,10 +558,13 @@ pandas==3.0.1 # via # datasets # evaluate +partial-json-parser==0.2.1.1.post7 + # via -r requirements/test/../common.txt pathvalidate==3.3.1 # via pytablewriter pillow==12.1.1 # via + # -r requirements/test/../common.txt # imageio # mistral-common # scikit-image @@ -410,16 +587,37 @@ portalocker==3.2.0 # via sacrebleu pqdm==0.2.0 # via -r requirements/test/xpu.in +prometheus-client==0.25.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # prometheus-fastapi-instrumentator +prometheus-fastapi-instrumentator==8.0.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt propcache==0.4.1 # via # aiohttp # yarl +protobuf==7.35.1 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # googleapis-common-protos + # opentelemetry-proto psutil==7.2.2 - # via accelerate + # via + # -r requirements/test/../common.txt + # accelerate py==1.11.0 # via pytest-forked +py-cpuinfo==9.0.0 + # via -r requirements/test/../common.txt pyarrow==23.0.1 # via datasets +pybase64==1.4.3 + # via -r requirements/test/../common.txt pycountry==26.2.16 # via pydantic-extra-types pycparser==3.0 @@ -429,23 +627,41 @@ pycryptodomex==3.23.0 pydantic==2.12.5 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # albumentations + # anthropic + # compressed-tensors # fastapi + # fastapi-cloud-cli # gpt-oss + # lm-format-enforcer + # mcp # mistral-common + # model-hosting-container-standards # mteb + # openai # openai-harmony # pydantic-extra-types + # pydantic-settings + # xgrammar pydantic-core==2.41.5 # via pydantic pydantic-extra-types==2.11.1 - # via mistral-common + # via + # fastapi + # mistral-common +pydantic-settings==2.14.2 + # via + # fastapi + # mcp pyelftools==0.32 # via triton-xpu pygments==2.20.0 # via # pytest # rich +pyjwt==2.13.0 + # via mcp pyrate-limiter==4.1.0 # via schemathesis pystemmer==3.0.0 @@ -480,19 +696,36 @@ python-dateutil==2.9.0.post0 # via # pandas # typepy +python-dotenv==1.2.2 + # via + # pydantic-settings + # uvicorn +python-json-logger==4.1.0 + # via -r requirements/test/../common.txt +python-multipart==0.0.32 + # via + # fastapi + # mcp pytrec-eval-terrier==0.5.10 # via mteb pytz==2026.1.post1 # via typepy pyyaml==6.0.3 # via + # -r requirements/test/../common.txt # accelerate # albumentations # datasets # huggingface-hub + # lm-format-enforcer # schemathesis # timm # transformers + # uvicorn +pyzmq==27.1.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt rapidfuzz==3.12.1 # via # -r requirements/test/xpu.in @@ -503,6 +736,7 @@ referencing==0.37.0 # jsonschema-specifications regex==2026.3.32 # via + # -r requirements/test/../common.txt # nltk # sacrebleu # tiktoken @@ -510,6 +744,7 @@ regex==2026.3.32 requests==2.33.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # datasets # docker # evaluate @@ -518,6 +753,7 @@ requests==2.33.1 # mistral-common # modelscope # mteb + # opentelemetry-exporter-otlp-proto-http # pooch # schemathesis # starlette-testclient @@ -525,8 +761,15 @@ requests==2.33.1 rich==14.3.3 # via # mteb + # rich-toolkit # schemathesis # typer +rich-toolkit==0.20.1 + # via + # fastapi-cli + # fastapi-cloud-cli +rignore==0.7.6 + # via fastapi-cloud-cli rouge-score==0.1.2 # via lm-eval rpds-py==0.30.0 @@ -538,6 +781,7 @@ sacrebleu==2.6.0 safetensors==0.7.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # accelerate # timm # transformers @@ -564,10 +808,18 @@ scipy==1.17.1 # sentence-transformers sentence-transformers==5.3.0 # via mteb +sentencepiece==0.2.1 + # via -r requirements/test/../common.txt +sentry-sdk==2.63.0 + # via fastapi-cloud-cli +setproctitle==1.3.7 + # via -r requirements/test/../common.txt setuptools==80.10.2 # via # -c requirements/common.txt # -c requirements/xpu.txt + # -r requirements/test/../common.txt + # model-hosting-container-standards # modelscope # pytablewriter # torch @@ -576,9 +828,14 @@ shellingham==1.5.4 six==1.17.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # junit-xml # python-dateutil # rouge-score +sniffio==1.3.1 + # via + # anthropic + # openai sortedcontainers==2.4.0 # via hypothesis soundfile==0.13.1 @@ -593,15 +850,24 @@ soxr==0.5.0.post1 # mistral-common sqlitedict==2.1.0 # via lm-eval +sse-starlette==3.4.5 + # via mcp starlette==1.3.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # fastapi + # mcp + # model-hosting-container-standards + # prometheus-fastapi-instrumentator + # sse-starlette # starlette-testclient starlette-testclient==0.4.1 # via schemathesis structlog==25.5.0 # via gpt-oss +supervisor==4.3.0 + # via model-hosting-container-standards sympy==1.14.0 # via torch tabledata==1.3.4 @@ -636,6 +902,7 @@ tifffile==2026.3.3 tiktoken==0.12.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss # lm-eval # mistral-common @@ -644,19 +911,23 @@ timm==1.0.17 tokenizers==0.22.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # transformers torch==2.12.0+xpu # via # -c requirements/xpu.txt # accelerate + # compressed-tensors # mteb # sentence-transformers # timm # torchvision + # xgrammar torchvision==0.27.0+xpu # via timm tqdm==4.67.3 # via + # -r requirements/test/../common.txt # datasets # evaluate # huggingface-hub @@ -664,13 +935,19 @@ tqdm==4.67.3 # modelscope # mteb # nltk + # openai # pqdm # sentence-transformers # transformers transformers==5.5.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # compressed-tensors # sentence-transformers + # xgrammar +triton==3.7.1 + # via xgrammar triton-xpu==3.7.1 # via torch typepy==1.3.4 @@ -680,36 +957,53 @@ typepy==1.3.4 # tabledata typer==0.24.1 # via + # fastapi-cli + # fastapi-cloud-cli # huggingface-hub # transformers typing-extensions==4.15.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # aiosignal # albumentations + # anthropic # anyio + # apache-tvm-ffi # chz # fastapi + # grpcio # huggingface-hub # librosa # lm-eval + # mcp # mistral-common # mteb + # openai + # opentelemetry-api + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions # pqdm # pydantic # pydantic-core # pydantic-extra-types # pytest-asyncio # referencing + # rich-toolkit # schemathesis # sentence-transformers # starlette # torch # typing-inspection + # xgrammar typing-inspection==0.4.2 # via # fastapi + # mcp # pydantic + # pydantic-settings umf==1.0.3 # via # intel-cmplr-lib-ur @@ -720,12 +1014,30 @@ urllib3==2.6.3 # docker # modelscope # requests + # sentry-sdk uvicorn==0.42.0 - # via gpt-oss + # via + # fastapi + # fastapi-cli + # fastapi-cloud-cli + # gpt-oss + # mcp +uvloop==0.22.1 + # via uvicorn +watchfiles==1.2.0 + # via + # -r requirements/test/../common.txt + # uvicorn +websockets==16.0 + # via uvicorn werkzeug==3.1.7 # via schemathesis word2number==1.1 # via lm-eval +xgrammar==0.2.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt xxhash==3.6.0 # via # datasets diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index bd84b2cbbfa..63c37a67554 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -102,45 +102,36 @@ def test_get_model_structural_tag_supports_vllm_hermes( ) assert isinstance(tag, StructuralTag) - assert tag.model_dump() == { - "type": "structural_tag", - "format": { - "type": "tags_with_separator", - "tags": [ - { - "type": "tag", - "begin": '\n{"name": "get_weather", "arguments": ', - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"], - }, - "style": "json", - }, - "end": "}\n", - }, - { - "type": "tag", - "begin": '{"name": "get_weather", "arguments": ', - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"], - }, - "style": "json", - }, - "end": "}", - }, - ], - "separator": "", - "at_least_one": True, - "stop_after_first": False, - }, + + # Assert the semantically meaningful structure rather than the full + # model_dump(), which gains version-specific keys across xgrammar releases + # (e.g. "any_order" was added to json_schema content in 0.2.3). + dump = tag.model_dump() + assert dump["type"] == "structural_tag" + + fmt = dump["format"] + assert fmt["type"] == "tags_with_separator" + assert fmt["separator"] == "" + assert fmt["at_least_one"] is True + assert fmt["stop_after_first"] is False + + expected_schema = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], } + expected_tags = [ + ('\n{"name": "get_weather", "arguments": ', "}\n"), + ('{"name": "get_weather", "arguments": ', "}"), + ] + assert len(fmt["tags"]) == len(expected_tags) + for tag_dump, (begin, end) in zip(fmt["tags"], expected_tags): + assert tag_dump["type"] == "tag" + assert tag_dump["begin"] == begin + assert tag_dump["end"] == end + content = tag_dump["content"] + assert content["type"] == "json_schema" + assert content["json_schema"] == expected_schema def test_hermes_required_tool_calls_use_empty_separator(): From 8fc1b2d046f4a991b5c24bb5470bf45efcfe9d01 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 29 Jun 2026 17:23:34 -0400 Subject: [PATCH 133/138] Fix FA4 dynamic_causal for full attention layers (#46659) Signed-off-by: Matthew Bonanni --- vllm/v1/attention/backends/flash_attn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index df209794352..2eed8190565 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -906,7 +906,10 @@ class FlashAttentionImpl(AttentionImpl): f"FA{self.vllm_flash_attn_version}" ) dynamic_causal = causal - causal = False + has_window = ( + sliding_window_size is not None and sliding_window_size[1] >= 0 + ) + causal = not has_window flash_attn_varlen_func( q=query[:num_actual_tokens], From ebcf511ec3c291eae63c38f5f431c07229e2406d Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Mon, 29 Jun 2026 16:24:08 -0500 Subject: [PATCH 134/138] [ROCm][CI] Soft Fail `Spec Decode Ngram + Suffix` and `Entrypoints Integration (LLM)` AMD Mirrors (#47067) Signed-off-by: Micah Williamson --- .buildkite/test_areas/entrypoints.yaml | 2 ++ .buildkite/test_areas/spec_decode.yaml | 2 ++ 2 files changed, 4 insertions(+) diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index d95b7e0d008..5ef88d4b97b 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -29,6 +29,8 @@ steps: mirror: amd: device: mi325_1 + # TODO(akaratza): Test after Torch >= 2.12 bump + soft_fail: true depends_on: - image-build-amd diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 6e532eddc71..671638f6f64 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -94,6 +94,8 @@ steps: amd: device: mi325_1 timeout_in_minutes: 65 + # TODO(akaratza): Test after Torch >= 2.12 bump + soft_fail: true depends_on: - image-build-amd source_file_dependencies: From 4eb227992aa2231ad538b8c90bc8191397ba3697 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 29 Jun 2026 16:26:41 -0500 Subject: [PATCH 135/138] [ROCm][CI] Make memory sampling less racy in tests and sleep mode (#45490) Signed-off-by: Andreas Karatzas Signed-off-by: Codex Co-authored-by: Codex --- tests/utils.py | 50 +++++++++++++++++++++++++++++++++--- vllm/v1/worker/gpu_worker.py | 16 +++++++++--- 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/tests/utils.py b/tests/utils.py index 2acb9716302..07cda56ce0e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1501,6 +1501,9 @@ def wait_for_gpu_memory_to_clear( threshold_bytes: int | dict[int, int] | None = None, threshold_ratio: float | dict[int, float] | None = None, timeout_s: float = 120, + stable_duration_s: float = 0, + stable_tolerance_bytes: int = 512 * 1024**2, + poll_interval_s: float = 5, ) -> None: assert threshold_bytes is not None or threshold_ratio is not None devices = get_physical_device_indices(devices) @@ -1528,8 +1531,13 @@ def wait_for_gpu_memory_to_clear( # Use nvml instead of pytorch to reduce measurement error from torch cuda # context. start_time = time.time() + stable_since: float | None = None + stable_used_bytes: dict[int, int] | None = None while True: output_raw = record_gpu_memory_usage_stats(devices=devices) + used_bytes_by_device = { + device: int(gb_used * 2**30) for device, (gb_used, _) in output_raw.items() + } output = { device: f"{gb_used:.02f}/{gb_total:.02f}" for device, (gb_used, gb_total) in output_raw.items() @@ -1577,15 +1585,45 @@ def wait_for_gpu_memory_to_clear( dur_s = time.time() - start_time if all_free: - print(f"Done waiting for free GPU memory on ({threshold=}) {dur_s=:.02f}") - break + if stable_duration_s <= 0: + print( + f"Done waiting for free GPU memory on devices {devices=} " + f"({threshold=}) {dur_s=:.02f}" + ) + break + + now = time.time() + if stable_used_bytes is None: + stable_since = now + stable_used_bytes = used_bytes_by_device + else: + memory_changed = any( + abs(used_bytes_by_device[device] - stable_used_bytes[device]) + > stable_tolerance_bytes + for device in devices + ) + if memory_changed: + stable_since = now + stable_used_bytes = used_bytes_by_device + elif ( + stable_since is not None and now - stable_since >= stable_duration_s + ): + print( + f"Done waiting for stable free GPU memory on devices " + f"{devices=} ({threshold=}) {dur_s=:.02f}" + ) + break + else: + stable_since = None + stable_used_bytes = None if dur_s >= timeout_s: raise ValueError( - f"Memory of devices not free after {dur_s=:.02f} ({threshold=})" + f"Memory of devices {devices=} not free after " + f"{dur_s=:.02f} ({threshold=})" ) - time.sleep(5) + time.sleep(poll_interval_s) def wait_for_rocm_memory_to_settle( @@ -1606,11 +1644,15 @@ def wait_for_rocm_memory_to_settle( num_gpus = current_platform.device_count() if num_gpus == 0: return + if threshold_ratio is None: + threshold_ratio = 0.1 wait_for_gpu_memory_to_clear( devices=list(range(num_gpus)), threshold_ratio=threshold_ratio, timeout_s=timeout_s, + stable_duration_s=2.0, + poll_interval_s=1.0, ) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 9afc2352528..07c3615edcb 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -4,6 +4,7 @@ import gc import os +import time from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta @@ -170,7 +171,8 @@ class Worker(WorkerBase): self._pp_send_work: list[Handle] = [] def sleep(self, level: int = 1) -> None: - free_bytes_before_sleep = torch.cuda.mem_get_info()[0] + torch.accelerator.synchronize() + free_bytes_before_sleep = current_platform.mem_get_info()[0] # Save the buffers before level 2 sleep if level == 2: @@ -181,8 +183,16 @@ class Worker(WorkerBase): allocator = get_mem_allocator_instance() allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) - free_bytes_after_sleep, total = torch.cuda.mem_get_info() - freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep + + torch.accelerator.synchronize() + deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0) + while True: + free_bytes_after_sleep, total = current_platform.mem_get_info() + freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep + if freed_bytes >= 0 or time.monotonic() >= deadline: + break + time.sleep(0.1) + used_bytes = total - free_bytes_after_sleep assert freed_bytes >= 0, "Memory usage increased after sleeping." logger.info( From 53f7553f099c2cbb88d2161959fc49dd71c8205b Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 29 Jun 2026 16:28:02 -0500 Subject: [PATCH 136/138] [ROCm][DeepEP] Stabilize high-throughput DBO for DP+EP (#46990) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: Tyler Michael Smith --- vllm/config/vllm.py | 21 +++++++++++++++++-- .../fused_moe/prepare_finalize/deepep_ht.py | 13 ++++++++++++ vllm/v1/worker/gpu_ubatch_wrapper.py | 10 +++++++++ 3 files changed, 42 insertions(+), 2 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b093b1788a9..b36c02a48ef 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -931,16 +931,28 @@ class VllmConfig: model_type, ) + from vllm.platforms import current_platform from vllm.v1.executor.abstract import Executor executor_backend = self.parallel_config.distributed_executor_backend executor_class = Executor.get_class(self) executor_supports_async_sched = executor_class.supports_async_scheduling() + uses_rocm_deepep_ht_dbo = ( + current_platform.is_rocm() + and self.parallel_config.enable_dbo + and self.parallel_config.all2all_backend == "deepep_high_throughput" + ) if self.scheduler_config.async_scheduling: # Async scheduling explicitly enabled, hard fail any incompatibilities. # Currently, async scheduling only support eagle speculative # decoding. + if uses_rocm_deepep_ht_dbo: + raise ValueError( + "Async scheduling is not compatible with ROCm DeepEP " + "high-throughput DBO. Please use --no-async-scheduling or " + "select a different all2all backend." + ) if self.speculative_config is not None: if ( self.speculative_config.method not in get_args(EagleModelTypes) @@ -1000,6 +1012,13 @@ class VllmConfig: executor_backend, ) self.scheduler_config.async_scheduling = False + elif uses_rocm_deepep_ht_dbo: + logger.warning_once( + "Async scheduling is disabled for ROCm DeepEP " + "high-throughput DBO because that combination can corrupt " + "DP+EP generation accuracy." + ) + self.scheduler_config.async_scheduling = False else: self.scheduler_config.async_scheduling = True @@ -1044,8 +1063,6 @@ class VllmConfig: "VLLM_WORKER_MULTIPROC_METHOD set to spawn" ) - from vllm.platforms import current_platform - if ( self.model_config is not None and self.scheduler_config.enable_chunked_prefill diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py index 45f9e815ac8..f30e0bea3c3 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceDelegate, ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.platforms import current_platform from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, @@ -59,6 +60,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): self.dp_size = dp_size self.rank_expert_offset = rank_expert_offset self.async_prepare = True + self.sync_dbo_comm = current_platform.is_rocm() # The dispatch function returns a handle that the combine function # requires. Under DBO microbatching we must track one handle per @@ -68,6 +70,13 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164 self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160] + def _sync_dbo_comm_if_needed(self) -> None: + if self.sync_dbo_comm and dbo_enabled(): + # ROCm DeepEP HT dispatch/combine reuse Buffer-owned communication + # workspace. Do not let the next DBO ubatch reuse that workspace + # before this ubatch's HT kernel has completed. + torch.cuda.current_stream().synchronize() + def num_dispatchers(self) -> int: return self.num_dispatchers_ @@ -161,6 +170,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): allocate_on_comm_stream=False, ) + self._sync_dbo_comm_if_needed() + # record the handle for this ubatch a2a_idx = dbo_current_ubatch_id() self.handles[a2a_idx] = handle @@ -375,6 +386,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): allocate_on_comm_stream=False, ) + self._sync_dbo_comm_if_needed() + dbo_switch_to_compute() if do_async: diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 657fc826734..76fa12b4121 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -154,6 +154,16 @@ class UBatchWrapper: @staticmethod def _create_sm_control_context(vllm_config: VllmConfig): comm_sms: int = envs.VLLM_DBO_COMM_SMS + rocm_deepep_ht_dbo = ( + current_platform.is_rocm() + and vllm_config.parallel_config.enable_dbo + and vllm_config.parallel_config.all2all_backend == "deepep_high_throughput" + ) + if rocm_deepep_ht_dbo: + # On ROCm, reserving CUs for DeepEP HT communication under DBO + # corrupts DP+EP generation accuracy. Keep the backend active, but + # leave all CUs visible to the compute and communication kernels. + comm_sms = 0 set_comm_sms = lambda sms: None if vllm_config.parallel_config.enable_expert_parallel: From c3734e8334ba124b722676e745084b8f4f86420b Mon Sep 17 00:00:00 2001 From: peizhang56 Date: Mon, 29 Jun 2026 14:29:47 -0700 Subject: [PATCH 137/138] [CI][Bugfix] Add cohere_melody to ROCm test requirements (#47072) Signed-off-by: pei.zhang Co-authored-by: Claude --- requirements/test/rocm.in | 1 + requirements/test/rocm.txt | 2 ++ 2 files changed, 3 insertions(+) diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index dc7f03c64f6..5afa6fcec92 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -70,6 +70,7 @@ gpt-oss>=0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank>=1.18.7 # required for fireredasr2 test +cohere_melody>=0.9.0 # required for cohere command reasoning parser test # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. # Older versions are in conflict with terratorch requirements. diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index b191705e0da..55cac6f5243 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -130,6 +130,8 @@ cloudpickle==3.1.2 # via # -r requirements/test/../common.txt # tilelang +cohere-melody==0.9.0 + # via -r requirements/test/rocm.in colorama==0.4.6 # via # perceptron From 8632c884dc440e231c8b7aef65a8795b80fe6676 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 29 Jun 2026 16:34:05 -0500 Subject: [PATCH 138/138] [ROCm][CI] Use spawn around the threaded OTLP test (#47003) Signed-off-by: Andreas Karatzas --- tests/v1/tracing/test_tracing.py | 121 ++++++++++++++++++------------- 1 file changed, 70 insertions(+), 51 deletions(-) diff --git a/tests/v1/tracing/test_tracing.py b/tests/v1/tracing/test_tracing.py index 2b450a6299c..1b7b243c9dc 100644 --- a/tests/v1/tracing/test_tracing.py +++ b/tests/v1/tracing/test_tracing.py @@ -7,6 +7,8 @@ import time from opentelemetry.sdk.environment_variables import OTEL_EXPORTER_OTLP_TRACES_INSECURE from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform from vllm.tracing import SpanAttributes # Import shared fixtures from the tracing conftest @@ -23,6 +25,11 @@ def test_traces( ): with monkeypatch.context() as m: m.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true") + if current_platform.is_rocm(): + # The fake OTLP server starts gRPC worker threads before the engine + # core is launched. On ROCm CI, forking while those threads are + # active can segfault in gRPC during engine startup or teardown. + m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") sampling_params = SamplingParams( temperature=0.01, @@ -30,58 +37,70 @@ def test_traces( max_tokens=256, ) model = "facebook/opt-125m" - llm = LLM( - model=model, - otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS, - gpu_memory_utilization=0.3, - disable_log_stats=False, - ) - prompts = ["This is a short prompt"] - outputs = llm.generate(prompts, sampling_params=sampling_params) - print(f"test_traces outputs is : {outputs}") + llm = None + try: + llm = LLM( + model=model, + otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS, + gpu_memory_utilization=0.3, + disable_log_stats=False, + ) + prompts = ["This is a short prompt"] + outputs = llm.generate(prompts, sampling_params=sampling_params) + print(f"test_traces outputs is : {outputs}") - # Wait for the "llm_request" span to be exported. - # The BatchSpanProcessor batches spans and exports them periodically, - # so we need to wait specifically for the llm_request span to appear. - timeout = 15 - deadline = time.time() + timeout - llm_request_spans = [] - while time.time() < deadline: - all_spans = trace_service.get_all_spans() - llm_request_spans = [s for s in all_spans if s["name"] == "llm_request"] - if llm_request_spans: - break - time.sleep(0.5) + # Wait for the "llm_request" span to be exported. + # The BatchSpanProcessor batches spans and exports them periodically, + # so we need to wait specifically for the llm_request span to appear. + timeout = 15 + deadline = time.time() + timeout + llm_request_spans = [] + while time.time() < deadline: + all_spans = trace_service.get_all_spans() + llm_request_spans = [s for s in all_spans if s["name"] == "llm_request"] + if llm_request_spans: + break + time.sleep(0.5) - assert len(llm_request_spans) == 1, ( - f"Expected exactly 1 'llm_request' span, but got {len(llm_request_spans)}. " - f"All span names: {[s['name'] for s in all_spans]}" - ) + assert len(llm_request_spans) == 1, ( + f"Expected exactly 1 'llm_request' span, but got " + f"{len(llm_request_spans)}. " + f"All span names: {[s['name'] for s in all_spans]}" + ) - attributes = llm_request_spans[0]["attributes"] - # assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model - assert attributes.get(SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id - assert ( - attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE) - == sampling_params.temperature - ) - assert ( - attributes.get(SpanAttributes.GEN_AI_REQUEST_TOP_P) == sampling_params.top_p - ) - assert ( - attributes.get(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS) - == sampling_params.max_tokens - ) - assert attributes.get(SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n - assert attributes.get(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len( - outputs[0].prompt_token_ids - ) - completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs) - assert ( - attributes.get(SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) - == completion_tokens - ) + attributes = llm_request_spans[0]["attributes"] + # assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_ID) + == outputs[0].request_id + ) + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE) + == sampling_params.temperature + ) + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_TOP_P) + == sampling_params.top_p + ) + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS) + == sampling_params.max_tokens + ) + assert attributes.get(SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n + assert attributes.get(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len( + outputs[0].prompt_token_ids + ) + completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs) + assert ( + attributes.get(SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) + == completion_tokens + ) - assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE) > 0 - assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) > 0 - assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) > 0 + assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE) > 0 + assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) > 0 + assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) > 0 + finally: + if llm is not None: + shutdown_timeout = 60.0 if current_platform.is_rocm() else 5.0 + llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout) + cleanup_dist_env_and_memory()