diff --git a/.buildkite/scripts/annotate-release.sh b/.buildkite/scripts/annotate-release.sh
index 56bb5cedaa0a9..df805e0850806 100755
--- a/.buildkite/scripts/annotate-release.sh
+++ b/.buildkite/scripts/annotate-release.sh
@@ -23,8 +23,8 @@ To download the wheel (by version):
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu126/vllm-${RELEASE_VERSION}+cu126-cp38-abi3-manylinux1_x86_64.whl .
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu129/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu130/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux1_x86_64.whl .
\`\`\`
To download and upload the image:
@@ -45,9 +45,10 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker push vllm/vllm-openai:latest-aarch64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 --amend
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 --amend
+docker manifest rm vllm/vllm-openai:latest
+docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest push vllm/vllm-openai:latest
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
\`\`\`
-EOF
\ No newline at end of file
+EOF
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
new file mode 100755
index 0000000000000..d0036f24c8d04
--- /dev/null
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
@@ -0,0 +1,64 @@
+#!/bin/bash
+
+# This script build the CPU docker image and run the offline inference inside the container.
+# It serves a sanity check for compilation and basic model usage.
+set -ex
+
+# allow to bind to different cores
+CORE_RANGE=${CORE_RANGE:-0-16}
+OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-16}
+NUMA_NODE=${NUMA_NODE:-0}
+
+export CMAKE_BUILD_PARALLEL_LEVEL=32
+
+# Setup cleanup
+remove_docker_container() {
+ set -e;
+ docker rm -f cpu-test-"$NUMA_NODE" || true;
+}
+trap remove_docker_container EXIT
+remove_docker_container
+
+# Try building the docker image
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
+
+# Run the image, setting --shm-size=4g for tensor parallel.
+docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
+
+function cpu_tests() {
+ set -e
+ export NUMA_NODE=$2
+
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
+ set -e
+ pip list"
+
+ # offline inference
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
+ set -e
+ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
+
+ # Run kernel tests
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
+ set -e
+ pytest -x -v -s tests/kernels/test_onednn.py
+ pytest -x -v -s tests/kernels/attention/test_cpu_attn.py"
+
+ # basic online serving
+ docker exec cpu-test-"$NUMA_NODE" bash -c '
+ set -e
+ VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve meta-llama/Llama-3.2-3B-Instruct --max-model-len 2048 &
+ server_pid=$!
+ timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
+ vllm bench serve \
+ --backend vllm \
+ --dataset-name random \
+ --model meta-llama/Llama-3.2-3B-Instruct \
+ --num-prompts 20 \
+ --endpoint /v1/completions
+ kill -s SIGTERM $server_pid &'
+}
+
+# All of CPU tests are expected to be finished less than 40 mins.
+export -f cpu_tests
+timeout 2h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
index 39ea180173081..3728f73fa2a36 100755
--- a/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
@@ -25,20 +25,22 @@ function cpu_tests() {
# offline inference
podman exec -it "$container_id" bash -c "
+ export TORCH_COMPILE_DISABLE=1
set -xve
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> $HOME/test_basic.log
# Run basic model test
podman exec -it "$container_id" bash -c "
+ export TORCH_COMPILE_DISABLE=1
set -evx
pip install pytest pytest-asyncio einops peft Pillow soundfile transformers_stream_generator matplotlib
- pip install sentence-transformers datamodel_code_generator
+ pip install sentence-transformers datamodel_code_generator tblib
# Note: disable Bart until supports V1
# pytest -v -s tests/models/language/generation/test_bart.py -m cpu_model
- pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-openai-community/gpt2]
- pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-facebook/opt-125m]
- pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-google/gemma-1.1-2b-it]
+ pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-openai-community/gpt2]
+ pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-facebook/opt-125m]
+ pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-google/gemma-1.1-2b-it]
pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
# TODO: Below test case tests/models/language/pooling/test_embedding.py::test_models[True-ssmits/Qwen2-7B-Instruct-embed-base] fails on ppc64le. Disabling it for time being.
# pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> $HOME/test_rest.log
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh
index 7479c43977d78..2267718f75ca5 100644
--- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh
@@ -73,12 +73,11 @@ function cpu_tests() {
pytest -x -s -v \
tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
- # Note: disable it until supports V1
- # Run AWQ test
- # docker exec cpu-test-"$NUMA_NODE" bash -c "
- # set -e
- # pytest -x -s -v \
- # tests/quantization/test_ipex_quant.py"
+ # Run AWQ/GPTQ test
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
+ set -e
+ pytest -x -s -v \
+ tests/quantization/test_cpu_wna16.py"
# Run multi-lora tests
docker exec cpu-test-"$NUMA_NODE" bash -c "
diff --git a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh
index 5302f524a0ae4..8106f50f18f66 100644
--- a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh
+++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh
@@ -17,7 +17,17 @@ wait_for_server() {
}
MODEL="deepseek-ai/DeepSeek-V2-lite"
-BACKENDS=("deepep_high_throughput" "deepep_low_latency")
+
+# Set BACKENDS based on platform
+if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
+ # ROCm platform
+ BACKENDS=("allgather_reducescatter")
+ # Disable MOE padding for ROCm since it is causing eplb to fail
+ export VLLM_ROCM_MOE_PADDING=0
+else
+ # Non-ROCm platform (CUDA/other)
+ BACKENDS=("deepep_high_throughput" "deepep_low_latency")
+fi
cleanup() {
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
diff --git a/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh
index a5135299297e2..0d06f53a183d0 100644
--- a/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh
+++ b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh
@@ -17,7 +17,16 @@ wait_for_server() {
}
MODEL="QWen/Qwen3-30B-A3B-FP8"
-BACKENDS=("deepep_high_throughput" "deepep_low_latency")
+# Set BACKENDS based on platform
+if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
+ # ROCm platform
+ BACKENDS=("allgather_reducescatter")
+ # Disable MOE padding for ROCm since it is causing eplb to fail
+ export VLLM_ROCM_MOE_PADDING=0
+else
+ # Non-ROCm platform (CUDA/other)
+ BACKENDS=("deepep_high_throughput" "deepep_low_latency")
+fi
cleanup() {
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index 2471b509a9fff..f098e23866eb3 100644
--- a/.buildkite/test-amd.yaml
+++ b/.buildkite/test-amd.yaml
@@ -61,7 +61,7 @@ steps:
- pytest -v -s -m 'not cpu_test' multimodal
- pytest -v -s utils_
-- label: Async Engine, Inputs, Utils, Worker Test (CPU) # 4 mins
+- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 4 mins
timeout_in_minutes: 10
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
@@ -73,6 +73,7 @@ steps:
- tests/multimodal
- tests/standalone_tests/lazy_imports.py
- tests/transformers_utils
+ - tests/config
no_gpu: true
commands:
- python3 standalone_tests/lazy_imports.py
@@ -80,6 +81,7 @@ steps:
- pytest -v -s test_outputs.py
- pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s transformers_utils
+ - pytest -v -s config
- label: Python-only Installation Test # 10min
timeout_in_minutes: 20
@@ -187,7 +189,7 @@ steps:
- tests/distributed/test_utils
- tests/distributed/test_pynccl
- tests/distributed/test_events
- - tests/compile/test_basic_correctness
+ - tests/compile/fullgraph/test_basic_correctness.py
- examples/offline_inference/rlhf.py
- examples/offline_inference/rlhf_colocate.py
- tests/examples/offline_inference/data_parallel.py
@@ -215,7 +217,7 @@ steps:
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
- pytest -v -s distributed/test_utils.py
- - pytest -v -s compile/test_basic_correctness.py
+ - pytest -v -s compile/fullgraph/test_basic_correctness.py
- pytest -v -s distributed/test_pynccl.py
- pytest -v -s distributed/test_events.py
- pytest -v -s distributed/test_symm_mem_allreduce.py
@@ -390,6 +392,15 @@ steps:
commands:
- pytest -v -s v1/attention
+- label: V1 Test attention (B200) # 10min
+ timeout_in_minutes: 30
+ gpu: b200
+ source_file_dependencies:
+ - vllm/v1/attention
+ - tests/v1/attention
+ commands:
+ - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
+
- label: V1 Test others (CPU) # 5 mins
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
@@ -493,17 +504,12 @@ steps:
- vllm/
- tests/compile
commands:
- - pytest -v -s compile/test_pass_manager.py
- - pytest -v -s compile/test_fusion.py
- - pytest -v -s compile/test_fusion_attn.py
- - pytest -v -s compile/test_functionalization.py
- - pytest -v -s compile/test_silu_mul_quant_fusion.py
- # - pytest -v -s compile/test_sequence_parallelism.py
- # - pytest -v -s compile/test_async_tp.py
- - pytest -v -s compile/test_fusion_all_reduce.py
- - pytest -v -s compile/test_decorator.py
- - pytest -v -s compile/test_noop_elimination.py
- - pytest -v -s compile/test_aot_compile.py
+ # Run unit tests defined directly under compile/,
+ # not including subdirectories, which are usually heavier
+ # tests covered elsewhere.
+ # Use `find` to launch multiple instances of pytest so that
+ # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
+ - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
- label: PyTorch Fullgraph Smoke Test # 15min
timeout_in_minutes: 30
@@ -515,9 +521,11 @@ steps:
- vllm/
- tests/compile
commands:
- - pytest -v -s compile/test_basic_correctness.py
- - pytest -v -s compile/test_multimodal_compile.py
- - pytest -v -s compile/piecewise/
+ # Run smoke tests under fullgraph directory, except test_full_graph.py
+ # as it is a heavy test that is covered in other steps.
+ # Use `find` to launch multiple instances of pytest so that
+ # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
+ - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;"
- label: PyTorch Fullgraph Test # 27min
timeout_in_minutes: 40
@@ -529,10 +537,10 @@ steps:
- vllm/
- tests/compile
commands:
- - pytest -v -s compile/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
+ - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
# Limit to no custom ops to reduce running time
# Wrap with quotes to escape yaml and avoid starting -k string with a -
- - "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and -quant_fp8'"
+ - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
- label: Cudagraph test
timeout_in_minutes: 20
@@ -697,7 +705,7 @@ steps:
- vllm/model_executor/models/whisper.py
commands: # LMEval
# Transcription WER check is skipped because encoder-decoder models are not supported on ROCm, see https://github.com/vllm-project/vllm/issues/27442
- - pytest -s entrypoints/openai/correctness/ --ignore entrypoints/openai/correctness/test_transcription_api_correctness.py
+ - pytest -s entrypoints/openai/correctness/
- label: OpenAI-Compatible Tool Use # 23 min
timeout_in_minutes: 35
@@ -746,6 +754,7 @@ steps:
torch_nightly: true
source_file_dependencies:
- vllm/model_executor/models/
+ - vllm/transformers_utils/
- tests/models/test_initialization.py
commands:
# Only when vLLM model source is modified - test initialization of a large
@@ -998,12 +1007,12 @@ steps:
optional: true
commands:
- pip install --upgrade git+https://github.com/huggingface/transformers
- - pytest -v -s tests/models/test_initialization.py
+ - pytest -v -s tests/models/test_initialization.py -k 'not (Gemma3 or ModernBert or Qwen2_5_VL or Qwen2_5vl or Qwen2VL or TransformersMultiModalEmbeddingModel or TransformersMultiModalForSequenceClassification or Ultravox or Phi4Multimodal or LlavaNextVideo or MiniCPMO or Lfm2Moe or PaliGemma or RobertaForSequenceClassification or Ovis2_5 or Fuyu or DeepseekOCR or KimiVL)'
- pytest -v -s tests/models/test_transformers.py
- - pytest -v -s tests/models/multimodal/processing/
- - pytest -v -s tests/models/multimodal/test_mapping.py
+ # - pytest -v -s tests/models/multimodal/processing/
+ - pytest -v -s tests/models/multimodal/test_mapping.py -k 'not (Gemma3 or Qwen2VL or Qwen2_5_VL)'
- python3 examples/offline_inference/basic/chat.py
- - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+ # - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
@@ -1048,7 +1057,7 @@ steps:
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
- pytest -v -s tests/kernels/moe/test_flashinfer.py
-- label: Blackwell Fusion Tests # 30 min
+- label: Blackwell Fusion and Compile Tests # 30 min
timeout_in_minutes: 40
working_dir: "/vllm-workspace/"
gpu: b200
@@ -1066,10 +1075,12 @@ steps:
- pytest -v -s tests/compile/test_fusion_attn.py
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
# this runner has 2 GPUs available even though num_gpus=2 is not set
- - pytest -v -s tests/compile/test_fusion_all_reduce.py
+ - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
# Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
# Wrap with quotes to escape yaml
- - "pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and Llama-3.1 and -quant_fp8 and -rms_norm'"
+ - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
+ # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
+ - pytest -v -s tests/compile/distributed/test_full_graph.py::test_fp8_kv_scale_compile
- label: Blackwell Fusion E2E Tests # 30 min
timeout_in_minutes: 40
@@ -1086,20 +1097,18 @@ steps:
- vllm/model_executor/layers/layernorm.py
- vllm/model_executor/layers/activation.py
- vllm/model_executor/layers/quantization/input_quant_fp8.py
- - tests/compile/test_fusions_e2e.py
- - tests/compile/test_full_graph.py
+ - tests/compile/distributed/test_fusions_e2e.py
+ - tests/compile/fullgraph/test_full_graph.py
commands:
- nvidia-smi
# Run all e2e fusion tests
- pytest -v -s tests/compile/test_fusions_e2e.py
- # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
- - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile
- label: ROCm GPT-OSS Eval
timeout_in_minutes: 60
working_dir: "/vllm-workspace/"
agent_pool: mi325_1
- mirror_hardwares: [amdproduction]
+ mirror_hardwares: [amdexperimental, amdproduction]
optional: true # run on nightlies
source_file_dependencies:
- tests/evals/gpt_oss
@@ -1198,7 +1207,7 @@ steps:
- vllm/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- - tests/compile/test_basic_correctness.py
+ - tests/compile/fullgraph/test_basic_correctness.py
- tests/compile/test_wrapper.py
- tests/distributed/
- tests/entrypoints/llm/test_collective_rpc.py
@@ -1211,7 +1220,7 @@ steps:
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
- pytest -v -s entrypoints/llm/test_collective_rpc.py
- - pytest -v -s ./compile/test_basic_correctness.py
+ - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
- pytest -v -s ./compile/test_wrapper.py
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
@@ -1311,7 +1320,10 @@ steps:
- pytest -v -s -x lora/test_llama_tp.py
- pytest -v -s -x lora/test_llm_with_multi_loras.py
- pytest -v -s -x lora/test_olmoe_tp.py
- - pytest -v -s -x lora/test_gptoss_tp.py
+
+ # Disabled for now because MXFP4 backend on non-cuda platform
+ # doesn't support LoRA yet
+ #- pytest -v -s -x lora/test_gptoss_tp.py
- label: Weight Loading Multiple GPU Test # 33min
@@ -1326,7 +1338,7 @@ steps:
- vllm/
- tests/weight_loading
commands:
- - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
+ - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
- label: Weight Loading Multiple GPU Test - Large Models # optional
mirror_hardwares: [amdexperimental]
@@ -1334,13 +1346,12 @@ steps:
# grade: Blocking
working_dir: "/vllm-workspace/tests"
num_gpus: 2
- gpu: a100
optional: true
source_file_dependencies:
- vllm/
- tests/weight_loading
commands:
- - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
+ - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
- label: NixlConnector PD accuracy tests (Distributed) # 30min
mirror_hardwares: [amdexperimental]
@@ -1417,10 +1428,12 @@ steps:
working_dir: "/vllm-workspace/"
num_gpus: 2
commands:
- - pytest -v -s tests/compile/test_async_tp.py
- - pytest -v -s tests/compile/test_sequence_parallelism.py
- - pytest -v -s tests/compile/test_fusion_all_reduce.py
- - pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
+ - pytest -v -s tests/compile/distributed/test_async_tp.py
+ - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
+ - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+ #- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
+ - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
+ - pytest -v -s tests/compile/distributed/test_sequence_parallel.py
- pytest -v -s tests/distributed/test_context_parallel.py
- CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048
- pytest -v -s tests/v1/distributed/test_dbo.py
diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml
index 4ac76aba67b9c..7a46e919f93bf 100644
--- a/.buildkite/test-pipeline.yaml
+++ b/.buildkite/test-pipeline.yaml
@@ -167,7 +167,7 @@ steps:
- tests/distributed/test_utils
- tests/distributed/test_pynccl
- tests/distributed/test_events
- - tests/compile/test_basic_correctness
+ - tests/compile/fullgraph/test_basic_correctness.py
- examples/offline_inference/rlhf.py
- examples/offline_inference/rlhf_colocate.py
- tests/examples/offline_inference/data_parallel.py
@@ -197,7 +197,7 @@ steps:
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
- pytest -v -s distributed/test_utils.py
- - pytest -v -s compile/test_basic_correctness.py
+ - pytest -v -s compile/fullgraph/test_basic_correctness.py
- pytest -v -s distributed/test_pynccl.py
- pytest -v -s distributed/test_events.py
- pytest -v -s distributed/test_symm_mem_allreduce.py
@@ -346,6 +346,18 @@ steps:
commands:
- pytest -v -s v1/attention
+- label: Batch Invariance Tests (H100) # 10min
+ timeout_in_minutes: 25
+ gpu: h100
+ source_file_dependencies:
+ - vllm/
+ - tests/v1/determinism/
+ commands:
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+ - pip install pytest-timeout pytest-forked
+ - pytest -v -s v1/determinism/test_batch_invariance.py
+ - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+
- label: V1 Test attention (B200) # 10min
timeout_in_minutes: 30
gpu: b200
@@ -445,18 +457,12 @@ steps:
- vllm/
- tests/compile
commands:
- - pytest -v -s compile/test_graph_partition.py
- - pytest -v -s compile/test_config.py
- - pytest -v -s compile/test_pass_manager.py
- - pytest -v -s compile/test_fusion.py
- - pytest -v -s compile/test_fusion_attn.py
- - pytest -v -s compile/test_functionalization.py
- - pytest -v -s compile/test_silu_mul_quant_fusion.py
- - pytest -v -s compile/test_fusion_all_reduce.py
- - pytest -v -s compile/test_decorator.py
- - pytest -v -s compile/test_noop_elimination.py
- - pytest -v -s compile/test_aot_compile.py
- - pytest -v -s compile/test_qk_norm_rope_fusion.py
+ # Run unit tests defined directly under compile/,
+ # not including subdirectories, which are usually heavier
+ # tests covered elsewhere.
+ # Use `find` to launch multiple instances of pytest so that
+ # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
+ - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
- label: PyTorch Fullgraph Smoke Test # 15min
timeout_in_minutes: 30
@@ -466,9 +472,11 @@ steps:
- vllm/
- tests/compile
commands:
- - pytest -v -s compile/test_basic_correctness.py
- - pytest -v -s compile/test_multimodal_compile.py
- - pytest -v -s compile/piecewise/
+ # Run smoke tests under fullgraph directory, except test_full_graph.py
+ # as it is a heavy test that is covered in other steps.
+ # Use `find` to launch multiple instances of pytest so that
+ # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
+ - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;"
- label: PyTorch Fullgraph Test # 27min
timeout_in_minutes: 40
@@ -479,10 +487,10 @@ steps:
- tests/compile
commands:
# fp8 kv scales not supported on sm89, tested on Blackwell instead
- - pytest -v -s compile/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
+ - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
# Limit to no custom ops to reduce running time
# Wrap with quotes to escape yaml and avoid starting -k string with a -
- - "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
+ - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
- label: Cudagraph test
timeout_in_minutes: 20
@@ -554,6 +562,25 @@ steps:
commands:
- pytest -v -s kernels/mamba
+- label: Kernels DeepGEMM Test (H100)
+ timeout_in_minutes: 45
+ gpu: h100
+ num_gpus: 1
+ source_file_dependencies:
+ - tools/install_deepgemm.sh
+ - vllm/utils/deep_gemm.py
+ - vllm/model_executor/layers/fused_moe
+ - vllm/model_executor/layers/quantization
+ - tests/kernels/quantization/test_block_fp8.py
+ - tests/kernels/moe/test_deepgemm.py
+ - tests/kernels/moe/test_batched_deepgemm.py
+ - tests/kernels/attention/test_deepgemm_attention.py
+ commands:
+ - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
+ - pytest -v -s kernels/moe/test_deepgemm.py
+ - pytest -v -s kernels/moe/test_batched_deepgemm.py
+ - pytest -v -s kernels/attention/test_deepgemm_attention.py
+
- label: Model Executor Test # 23min
timeout_in_minutes: 35
torch_nightly: true
@@ -664,6 +691,7 @@ steps:
torch_nightly: true
source_file_dependencies:
- vllm/model_executor/models/
+ - vllm/transformers_utils/
- tests/models/test_initialization.py
commands:
# Only when vLLM model source is modified - test initialization of a large
@@ -876,12 +904,12 @@ steps:
optional: true
commands:
- pip install --upgrade git+https://github.com/huggingface/transformers
- - pytest -v -s tests/models/test_initialization.py -k 'not (Gemma3 or ModernBert or Qwen2_5_VL or Qwen2_5vl or Qwen2VL or TransformersMultiModalEmbeddingModel or TransformersMultiModalForSequenceClassification or Ultravox or Phi4Multimodal or LlavaNextVideo or MiniCPMO or Lfm2Moe or PaliGemma or RobertaForSequenceClassification or Ovis2_5 or Fuyu or DeepseekOCR or KimiVL)'
+ - pytest -v -s tests/models/test_initialization.py -k 'not (Ultravox or Phi4Multimodal or MiniCPMO or Lfm2Moe or RobertaForSequenceClassification or Ovis2_5 or DeepseekOCR or KimiVL)'
- pytest -v -s tests/models/test_transformers.py
# - pytest -v -s tests/models/multimodal/processing/
- - pytest -v -s tests/models/multimodal/test_mapping.py -k 'not (Gemma3 or Qwen2VL or Qwen2_5_VL)'
+ - pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/offline_inference/basic/chat.py
- # - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+ - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
@@ -925,6 +953,7 @@ steps:
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
- pytest -v -s tests/kernels/moe/test_flashinfer.py
+ - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
- label: Blackwell Fusion and Compile Tests # 30 min
timeout_in_minutes: 40
@@ -934,22 +963,29 @@ steps:
- csrc/quantization/fp4/
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
- vllm/v1/attention/backends/flashinfer.py
+ - vllm/v1/worker/
+ - vllm/v1/cudagraph_dispatcher.py
- vllm/compilation/
# can affect pattern matching
- vllm/model_executor/layers/layernorm.py
- vllm/model_executor/layers/activation.py
- vllm/model_executor/layers/quantization/input_quant_fp8.py
+ - tests/compile/test_fusion_attn.py
+ - tests/compile/test_silu_mul_quant_fusion.py
+ - tests/compile/distributed/test_fusion_all_reduce.py
+ - tests/compile/distributed/test_fusions_e2e.py
+ - tests/compile/fullgraph/test_full_graph.py
commands:
- nvidia-smi
- pytest -v -s tests/compile/test_fusion_attn.py
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
# this runner has 2 GPUs available even though num_gpus=2 is not set
- - pytest -v -s tests/compile/test_fusion_all_reduce.py
+ - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
# Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
# Wrap with quotes to escape yaml
- - "pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
+ - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
# test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
- - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile
+ - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
- label: Blackwell Fusion E2E Tests # 30 min
timeout_in_minutes: 40
@@ -966,12 +1002,11 @@ steps:
- vllm/model_executor/layers/layernorm.py
- vllm/model_executor/layers/activation.py
- vllm/model_executor/layers/quantization/input_quant_fp8.py
- - tests/compile/test_fusions_e2e.py
- - tests/compile/test_full_graph.py
+ - tests/compile/distributed/test_fusions_e2e.py
commands:
- nvidia-smi
# Run all e2e fusion tests
- - pytest -v -s tests/compile/test_fusions_e2e.py
+ - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
- label: Blackwell GPT-OSS Eval
timeout_in_minutes: 60
@@ -1069,7 +1104,7 @@ steps:
- vllm/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- - tests/compile/test_basic_correctness.py
+ - tests/compile/fullgraph/test_basic_correctness.py
- tests/compile/test_wrapper.py
- tests/distributed/
- tests/entrypoints/llm/test_collective_rpc.py
@@ -1084,7 +1119,7 @@ steps:
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
- pytest -v -s entrypoints/llm/test_collective_rpc.py
- - pytest -v -s ./compile/test_basic_correctness.py
+ - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
- pytest -v -s ./compile/test_wrapper.py
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
@@ -1264,10 +1299,10 @@ steps:
working_dir: "/vllm-workspace/"
num_gpus: 2
commands:
- - pytest -v -s tests/compile/test_async_tp.py
- - pytest -v -s tests/compile/test_sequence_parallelism.py
- - pytest -v -s tests/compile/test_fusion_all_reduce.py
- - "pytest -v -s tests/compile/test_fusions_e2e.py -k 'not Llama-4'"
+ - pytest -v -s tests/compile/distributed/test_async_tp.py
+ - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
+ - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+ - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
- pytest -v -s tests/distributed/test_sequence_parallel.py
- pytest -v -s tests/distributed/test_context_parallel.py
- CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 6e178bb690c56..3247408e1163e 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -9,6 +9,7 @@
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
/vllm/model_executor/layers/mamba @tdoublep
/vllm/model_executor/model_loader @22quinn
+/vllm/model_executor/layers/batch_invariant.py @yewentao256
/vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
/vllm/vllm_flash_attn @LucasWilkinson
/vllm/lora @jeejeelee
@@ -35,6 +36,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/vllm/v1/kv_cache_interface.py @heheda12345
/vllm/v1/offloading @ApostaC
+# Model runner V2
+/vllm/v1/worker/gpu @WoosukKwon
+
# Test ownership
/.buildkite/lm-eval-harness @mgoin
/tests/distributed/test_multi_node_assignment.py @youkaichao
@@ -56,6 +60,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/tests/v1/kv_connector/nixl_integration @NickLucche
/tests/v1/kv_connector @ApostaC
/tests/v1/offloading @ApostaC
+/tests/v1/determinism @yewentao256
# Transformers modeling backend
/vllm/model_executor/models/transformers @hmellor
diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml
index 8d40aa587bf00..a183033c9adde 100644
--- a/.github/workflows/macos-smoke-test.yml
+++ b/.github/workflows/macos-smoke-test.yml
@@ -1,12 +1,15 @@
name: macOS Apple Silicon Smoke Test
on:
+ push:
+ branches:
+ - main
workflow_dispatch: # Manual trigger
jobs:
macos-m1-smoke-test:
runs-on: macos-latest
- timeout-minutes: 20
+ timeout-minutes: 30
steps:
- uses: actions/checkout@v4
@@ -19,28 +22,29 @@ jobs:
pyproject.toml
python-version: '3.12'
- - name: Install dependencies
+ - name: Create virtual environment
run: |
- uv pip install -r requirements/cpu-build.txt
- uv pip install -r requirements/cpu.txt
+ uv venv
+ echo "$GITHUB_WORKSPACE/.venv/bin" >> "$GITHUB_PATH"
- - name: Build vLLM
- run: uv pip install -v -e .
+ - name: Install dependencies and build vLLM
+ run: |
+ uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match
+ uv pip install -e .
env:
CMAKE_BUILD_PARALLEL_LEVEL: 4
- name: Verify installation
run: |
python -c "import vllm; print(f'vLLM version: {vllm.__version__}')"
- python -c "import torch; print(f'PyTorch: {torch.__version__}')"
- name: Smoke test vllm serve
- timeout-minutes: 10
run: |
# Start server in background
vllm serve Qwen/Qwen3-0.6B \
- --max-model-len=2048 \
+ --max-model-len=2K \
--load-format=dummy \
+ --hf-overrides '{"num_hidden_layers": 2}' \
--enforce-eager \
--port 8000 &
diff --git a/.gitignore b/.gitignore
index 50070d7898fe6..7cda86478664f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,9 @@
# vllm-flash-attn built from source
vllm/vllm_flash_attn/*
+# OpenAI triton kernels copied from source
+vllm/third_party/triton_kernels/*
+
# triton jit
.triton
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3a37040edbf1a..a4cf51d17e982 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -307,7 +307,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
- set(CUTLASS_REVISION "v4.2.1" CACHE STRING "CUTLASS revision to use")
+ set(CUTLASS_REVISION "v4.2.1")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
@@ -512,9 +512,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
# The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
# require CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
set(SRCS
@@ -619,9 +619,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
# FP4 Archs and flags
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
- cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
- cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;12.0a;12.1a" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
set(SRCS
@@ -695,7 +695,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
set(SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
@@ -741,9 +741,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
set(SRCS "csrc/quantization/w8a8/cutlass/moe/blockwise_scaled_group_mm_sm100.cu")
@@ -1030,6 +1030,11 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
WITH_SOABI)
endif()
+# For CUDA and HIP builds also build the triton_kernels external package.
+if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
+ include(cmake/external_projects/triton_kernels.cmake)
+endif()
+
# For CUDA we also build and ship some external projects.
if (VLLM_GPU_LANG STREQUAL "CUDA")
include(cmake/external_projects/flashmla.cmake)
diff --git a/benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py b/benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py
index 904f805349148..d072c03c440b2 100644
--- a/benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py
+++ b/benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py
@@ -5,11 +5,12 @@ import argparse
import asyncio
import logging
import os
+import time
+import uuid
+from urllib.parse import urlparse
import aiohttp
from quart import Quart, Response, make_response, request
-from rate_limiter import RateLimiter
-from request_queue import RequestQueue
# Configure logging
logging.basicConfig(level=logging.INFO)
@@ -24,26 +25,8 @@ def parse_args():
parser.add_argument(
"--timeout",
type=float,
- default=300,
- help="Timeout for backend service requests in seconds (default: 300)",
- )
- parser.add_argument(
- "--max-concurrent",
- type=int,
- default=100,
- help="Maximum concurrent requests to backend services (default: 100)",
- )
- parser.add_argument(
- "--queue-size",
- type=int,
- default=500,
- help="Maximum number of requests in the queue (default: 500)",
- )
- parser.add_argument(
- "--rate-limit",
- type=int,
- default=40,
- help="Maximum requests per second (default: 40)",
+ default=6 * 60 * 60,
+ help="Timeout for backend service requests in seconds (default: 21600)",
)
parser.add_argument(
"--port",
@@ -54,14 +37,32 @@ def parse_args():
parser.add_argument(
"--prefill-url",
type=str,
- default="http://localhost:8100/v1/completions",
- help="Prefill service endpoint URL",
+ default="http://localhost:8100",
+ help="Prefill service base URL (protocol + host[:port])",
)
parser.add_argument(
"--decode-url",
type=str,
- default="http://localhost:8200/v1/completions",
- help="Decode service endpoint URL",
+ default="http://localhost:8200",
+ help="Decode service base URL (protocol + host[:port])",
+ )
+ parser.add_argument(
+ "--kv-host",
+ type=str,
+ default="localhost",
+ help="Hostname or IP used by KV transfer (default: localhost)",
+ )
+ parser.add_argument(
+ "--prefill-kv-port",
+ type=int,
+ default=14579,
+ help="Prefill KV port (default: 14579)",
+ )
+ parser.add_argument(
+ "--decode-kv-port",
+ type=int,
+ default=14580,
+ help="Decode KV port (default: 14580)",
)
return parser.parse_args()
@@ -73,70 +74,129 @@ def main():
# Initialize configuration using command line parameters
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
- MAX_CONCURRENT_REQUESTS = args.max_concurrent
- REQUEST_QUEUE_SIZE = args.queue_size
- RATE_LIMIT = args.rate_limit
PREFILL_SERVICE_URL = args.prefill_url
DECODE_SERVICE_URL = args.decode_url
PORT = args.port
+ PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
+ DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
+
+ logger.info(
+ "Proxy resolved KV addresses -> prefill: %s, decode: %s",
+ PREFILL_KV_ADDR,
+ DECODE_KV_ADDR,
+ )
+
app = Quart(__name__)
- # Initialize the rate limiter and request queue
- rate_limiter = RateLimiter(RATE_LIMIT)
- request_queue = RequestQueue(MAX_CONCURRENT_REQUESTS, REQUEST_QUEUE_SIZE)
-
- # Attach the configuration object to the application instance
+ # Attach the configuration object to the application instance so helper
+ # coroutines can read the resolved backend URLs and timeouts without using
+ # globals.
app.config.update(
{
"AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
- "rate_limiter": rate_limiter,
- "request_queue": request_queue,
"PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
"DECODE_SERVICE_URL": DECODE_SERVICE_URL,
+ "PREFILL_KV_ADDR": PREFILL_KV_ADDR,
+ "DECODE_KV_ADDR": DECODE_KV_ADDR,
}
)
- # Start queue processing on app startup
- @app.before_serving
- async def startup():
- """Start request processing task when app starts serving"""
- asyncio.create_task(request_queue.process())
+ def _normalize_base_url(url: str) -> str:
+ """Remove any trailing slash so path joins behave predictably."""
+ return url.rstrip("/")
- async def forward_request(url, data):
- """Forward request to backend service with rate limiting and error handling"""
- headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
+ def _get_host_port(url: str) -> str:
+ """Return the hostname:port portion for logging and KV headers."""
+ parsed = urlparse(url)
+ host = parsed.hostname or "localhost"
+ port = parsed.port
+ if port is None:
+ port = 80 if parsed.scheme == "http" else 443
+ return f"{host}:{port}"
- # Use rate limiter as context manager
- async with (
- rate_limiter,
- aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
- ):
- try:
- async with session.post(
- url=url, json=data, headers=headers
- ) as response:
- if response.status == 200:
- # Stream response chunks
- async for chunk_bytes in response.content.iter_chunked(1024):
- yield chunk_bytes
- else:
- # Handle backend service errors
- error_text = await response.text()
- logger.error(
- "Backend service error: %s - %s",
- response.status,
- error_text,
- )
- yield b'{"error": "Backend service error"}'
- except aiohttp.ClientError as e:
- # Handle connection errors
- logger.error("Connection error to %s: %s", url, str(e))
- yield b'{"error": "Service unavailable"}'
- except asyncio.TimeoutError:
- # Handle timeout errors
- logger.error("Timeout connecting to %s", url)
- yield b'{"error": "Service timeout"}'
+ PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
+ DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
+ KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
+
+ def _build_headers(request_id: str) -> dict[str, str]:
+ """Construct the headers expected by vLLM's P2P disagg connector."""
+ headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
+ api_key = os.environ.get("OPENAI_API_KEY")
+ if api_key:
+ headers["Authorization"] = f"Bearer {api_key}"
+ return headers
+
+ async def _run_prefill(
+ request_path: str,
+ payload: dict,
+ headers: dict[str, str],
+ request_id: str,
+ ):
+ url = f"{PREFILL_BASE}{request_path}"
+ start_ts = time.perf_counter()
+ logger.info("[prefill] start request_id=%s url=%s", request_id, url)
+ try:
+ async with (
+ aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
+ session.post(url=url, json=payload, headers=headers) as resp,
+ ):
+ if resp.status != 200:
+ error_text = await resp.text()
+ raise RuntimeError(
+ f"Prefill backend error {resp.status}: {error_text}"
+ )
+ await resp.read()
+ logger.info(
+ "[prefill] done request_id=%s status=%s elapsed=%.2fs",
+ request_id,
+ resp.status,
+ time.perf_counter() - start_ts,
+ )
+ except asyncio.TimeoutError as exc:
+ raise RuntimeError(f"Prefill service timeout at {url}") from exc
+ except aiohttp.ClientError as exc:
+ raise RuntimeError(f"Prefill service unavailable at {url}") from exc
+
+ async def _stream_decode(
+ request_path: str,
+ payload: dict,
+ headers: dict[str, str],
+ request_id: str,
+ ):
+ url = f"{DECODE_BASE}{request_path}"
+ # Stream tokens from the decode service once the prefill stage has
+ # materialized KV caches on the target workers.
+ logger.info("[decode] start request_id=%s url=%s", request_id, url)
+ try:
+ async with (
+ aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
+ session.post(url=url, json=payload, headers=headers) as resp,
+ ):
+ if resp.status != 200:
+ error_text = await resp.text()
+ logger.error(
+ "Decode backend error %s - %s", resp.status, error_text
+ )
+ err_msg = (
+ '{"error": "Decode backend error ' + str(resp.status) + '"}'
+ )
+ yield err_msg.encode()
+ return
+ logger.info(
+ "[decode] streaming response request_id=%s status=%s",
+ request_id,
+ resp.status,
+ )
+ async for chunk_bytes in resp.content.iter_chunked(1024):
+ yield chunk_bytes
+ logger.info("[decode] finished streaming request_id=%s", request_id)
+ except asyncio.TimeoutError:
+ logger.error("Decode service timeout at %s", url)
+ yield b'{"error": "Decode service timeout"}'
+ except aiohttp.ClientError as exc:
+ logger.error("Decode service error at %s: %s", url, exc)
+ yield b'{"error": "Decode service unavailable"}'
async def process_request():
"""Process a single request through prefill and decode stages"""
@@ -146,13 +206,27 @@ def main():
# Create prefill request (max_tokens=1)
prefill_request = original_request_data.copy()
prefill_request["max_tokens"] = 1
+ if "max_completion_tokens" in prefill_request:
+ prefill_request["max_completion_tokens"] = 1
# Execute prefill stage
- async for _ in forward_request(PREFILL_SERVICE_URL, prefill_request):
- continue
+ # The request id encodes both KV socket addresses so the backend can
+ # shuttle tensors directly via NCCL once the prefill response
+ # completes.
+ request_id = (
+ f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
+ f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
+ )
+
+ headers = _build_headers(request_id)
+ await _run_prefill(request.path, prefill_request, headers, request_id)
# Execute decode stage and stream response
- generator = forward_request(DECODE_SERVICE_URL, original_request_data)
+ # Pass the unmodified user request so the decode phase can continue
+ # sampling with the already-populated KV cache.
+ generator = _stream_decode(
+ request.path, original_request_data, headers, request_id
+ )
response = await make_response(generator)
response.timeout = None # Disable timeout for streaming response
return response
@@ -168,23 +242,10 @@ def main():
@app.route("/v1/completions", methods=["POST"])
async def handle_request():
"""Handle incoming API requests with concurrency and rate limiting"""
- # Create task for request processing
- task = asyncio.create_task(process_request())
-
- # Enqueue request or reject if queue is full
- if not await request_queue.enqueue(task):
- return Response(
- response=b'{"error": "Server busy, try again later"}',
- status=503,
- content_type="application/json",
- )
-
try:
- # Return the response from the processing task
- return await task
+ return await process_request()
except asyncio.CancelledError:
- # Handle task cancellation (timeout or queue full)
- logger.warning("Request cancelled due to timeout or queue full")
+ logger.warning("Request cancelled")
return Response(
response=b'{"error": "Request cancelled"}',
status=503,
diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py
index 027f67ad4db69..e07d6c776bc00 100644
--- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py
+++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py
@@ -255,8 +255,8 @@ def bench_run(
torch.cuda.synchronize()
# Timing
- start_event = torch.cuda.Event(enable_timing=True)
- end_event = torch.cuda.Event(enable_timing=True)
+ start_event = torch.Event(enable_timing=True)
+ end_event = torch.Event(enable_timing=True)
latencies = []
for _ in range(num_iters):
diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py
index c99951aa27826..a1af0b8aec3d0 100644
--- a/benchmarks/kernels/benchmark_moe.py
+++ b/benchmarks/kernels/benchmark_moe.py
@@ -185,8 +185,8 @@ def benchmark_config(
graph.replay()
torch.cuda.synchronize()
- start_event = torch.cuda.Event(enable_timing=True)
- end_event = torch.cuda.Event(enable_timing=True)
+ start_event = torch.Event(enable_timing=True)
+ end_event = torch.Event(enable_timing=True)
latencies: list[float] = []
for i in range(num_iters):
diff --git a/benchmarks/kernels/benchmark_moe_permute_unpermute.py b/benchmarks/kernels/benchmark_moe_permute_unpermute.py
index efa5a7386027e..b8913a217c608 100644
--- a/benchmarks/kernels/benchmark_moe_permute_unpermute.py
+++ b/benchmarks/kernels/benchmark_moe_permute_unpermute.py
@@ -105,8 +105,8 @@ def benchmark_permute(
graph.replay()
torch.cuda.synchronize()
- start_event = torch.cuda.Event(enable_timing=True)
- end_event = torch.cuda.Event(enable_timing=True)
+ start_event = torch.Event(enable_timing=True)
+ end_event = torch.Event(enable_timing=True)
latencies: list[float] = []
for i in range(num_iters):
@@ -241,8 +241,8 @@ def benchmark_unpermute(
graph.replay()
torch.cuda.synchronize()
- start_event = torch.cuda.Event(enable_timing=True)
- end_event = torch.cuda.Event(enable_timing=True)
+ start_event = torch.Event(enable_timing=True)
+ end_event = torch.Event(enable_timing=True)
latencies: list[float] = []
for i in range(num_iters):
diff --git a/benchmarks/kernels/benchmark_mrope.py b/benchmarks/kernels/benchmark_mrope.py
index cb848d2bf579e..83bd91917508f 100644
--- a/benchmarks/kernels/benchmark_mrope.py
+++ b/benchmarks/kernels/benchmark_mrope.py
@@ -6,7 +6,7 @@
#
# The CSV file (named with current date/time) contains these columns:
# model_name, tp_size, num_tokens, num_heads, num_kv_heads, head_dim, max_position,
-# rope_theta, is_neox_style, rope_scaling, dtype, torch_mean, torch_median, torch_p99,
+# is_neox_style, rope_parameters, dtype, torch_mean, torch_median, torch_p99,
# torch_min, torch_max, triton_mean, triton_median, triton_p99, triton_min, triton_max,
# speedup
#
@@ -86,9 +86,8 @@ def benchmark_mrope(
num_heads: int,
num_kv_heads: int,
max_position: int = 8192,
- rope_theta: float = 10000,
is_neox_style: bool = True,
- rope_scaling: dict[str, Any] = None,
+ rope_parameters: dict[str, Any] | None = None,
dtype: torch.dtype = torch.bfloat16,
seed: int = 0,
warmup_iter: int = 10,
@@ -102,9 +101,8 @@ def benchmark_mrope(
head_size=head_dim,
rotary_dim=head_dim,
max_position=max_position,
- base=rope_theta,
is_neox_style=is_neox_style,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
dtype=dtype,
).to(device=device)
@@ -203,9 +201,8 @@ def benchmark_mrope(
num_kv_heads,
head_dim,
max_position,
- rope_theta,
is_neox_style,
- str(rope_scaling),
+ str(rope_parameters),
str(dtype).split(".")[-1],
torch_stats["mean"],
torch_stats["median"],
@@ -255,9 +252,8 @@ if __name__ == "__main__":
"num_kv_heads",
"head_dim",
"max_position",
- "rope_theta",
"is_neox_style",
- "rope_scaling",
+ "rope_parameters",
"dtype",
"torch_mean",
"torch_median",
@@ -303,7 +299,7 @@ if __name__ == "__main__":
q_size = num_heads * head_dim
kv_size = num_kv_heads * head_dim
is_neox_style = True
- rope_theta = config.rope_theta
+ rope_parameters = config.rope_parameters
max_position = config.max_position_embeddings
for num_tokens in num_tokens_list:
@@ -315,9 +311,8 @@ if __name__ == "__main__":
num_heads=num_heads,
num_kv_heads=num_kv_heads,
max_position=max_position,
- rope_theta=rope_theta,
is_neox_style=is_neox_style,
- rope_scaling=config.rope_scaling,
+ rope_parameters=rope_parameters,
dtype=getattr(torch, args.dtype),
seed=args.seed,
warmup_iter=args.warmup_iter,
diff --git a/benchmarks/kernels/benchmark_per_token_group_quant.py b/benchmarks/kernels/benchmark_per_token_group_quant.py
index bdc1eb733084e..eba4d510258b6 100644
--- a/benchmarks/kernels/benchmark_per_token_group_quant.py
+++ b/benchmarks/kernels/benchmark_per_token_group_quant.py
@@ -30,8 +30,8 @@ def _time_cuda(
fn()
torch.cuda.synchronize()
- start = torch.cuda.Event(enable_timing=True)
- end = torch.cuda.Event(enable_timing=True)
+ start = torch.Event(enable_timing=True)
+ end = torch.Event(enable_timing=True)
start.record()
for _ in range(bench_iters):
diff --git a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
index a5887aafd30d6..de01ff197eab7 100644
--- a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
+++ b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
@@ -253,8 +253,8 @@ def benchmark(
)
torch.cuda.synchronize()
- start_event = torch.cuda.Event(enable_timing=True)
- end_event = torch.cuda.Event(enable_timing=True)
+ start_event = torch.Event(enable_timing=True)
+ end_event = torch.Event(enable_timing=True)
# Benchmark
latencies: list[float] = []
diff --git a/benchmarks/kernels/benchmark_trtllm_decode_attention.py b/benchmarks/kernels/benchmark_trtllm_decode_attention.py
index 29ce18234dfa0..1d0d6fbb9a470 100644
--- a/benchmarks/kernels/benchmark_trtllm_decode_attention.py
+++ b/benchmarks/kernels/benchmark_trtllm_decode_attention.py
@@ -127,8 +127,8 @@ def benchmark_decode(
def time_fn(fn, warmup=10, trials=20):
torch.cuda.synchronize()
- start = torch.cuda.Event(enable_timing=True)
- end = torch.cuda.Event(enable_timing=True)
+ start = torch.Event(enable_timing=True)
+ end = torch.Event(enable_timing=True)
times = []
for i in range(warmup):
fn()
diff --git a/benchmarks/kernels/benchmark_trtllm_prefill_attention.py b/benchmarks/kernels/benchmark_trtllm_prefill_attention.py
index 2a25d03748112..84bde723abf7f 100644
--- a/benchmarks/kernels/benchmark_trtllm_prefill_attention.py
+++ b/benchmarks/kernels/benchmark_trtllm_prefill_attention.py
@@ -139,8 +139,8 @@ def benchmark_prefill(
def time_fn(fn, warmup=10, trials=20):
torch.cuda.synchronize()
- start = torch.cuda.Event(enable_timing=True)
- end = torch.cuda.Event(enable_timing=True)
+ start = torch.Event(enable_timing=True)
+ end = torch.Event(enable_timing=True)
times = []
for i in range(warmup):
fn()
diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py
index ab54f81985bc2..b52500c8c5217 100644
--- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py
+++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py
@@ -183,8 +183,8 @@ def benchmark_config(
run()
torch.cuda.synchronize()
- start_event = torch.cuda.Event(enable_timing=True)
- end_event = torch.cuda.Event(enable_timing=True)
+ start_event = torch.Event(enable_timing=True)
+ end_event = torch.Event(enable_timing=True)
latencies: list[float] = []
for i in range(num_iters):
diff --git a/benchmarks/multi_turn/README.md b/benchmarks/multi_turn/README.md
index f5b5c6c97d484..b0be1e3a69a66 100644
--- a/benchmarks/multi_turn/README.md
+++ b/benchmarks/multi_turn/README.md
@@ -55,6 +55,10 @@ output_num_chunks 166.0 99.01 11.80 79.00 90.00 98.00 108.75
----------------------------------------------------------------------------------------------------
```
+If you run with `--warmup-step`, the summary will also include `warmup_runtime_sec`
+and `total_runtime_incl_warmup_sec` (while `runtime_sec` continues to reflect the
+benchmark-only runtime so the reported throughput stays comparable).
+
### JSON configuration file for synthetic conversations generation
The input flag `--input-file` is used to determine the input conversations for the benchmark.
diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py
index 772d685ad90ff..e23f6b923f1b9 100644
--- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py
+++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py
@@ -1076,6 +1076,7 @@ def process_statistics(
verbose: bool,
gen_conv_args: GenConvArgs | None = None,
excel_output: bool = False,
+ warmup_runtime_sec: float | None = None,
) -> None:
if len(client_metrics) == 0:
logger.info("No samples to process")
@@ -1169,8 +1170,13 @@ def process_statistics(
# Convert milliseconds to seconds
runtime_sec = runtime_sec / 1000.0
requests_per_sec = float(len(df)) / runtime_sec
-
- params = {"runtime_sec": runtime_sec, "requests_per_sec": requests_per_sec}
+ params = {
+ "runtime_sec": runtime_sec,
+ "requests_per_sec": requests_per_sec,
+ }
+ if warmup_runtime_sec is not None:
+ params["warmup_runtime_sec"] = warmup_runtime_sec
+ params["total_runtime_incl_warmup_sec"] = runtime_sec + warmup_runtime_sec
# Generate a summary of relevant metrics (and drop irrelevant data)
df = df.drop(columns=exclude).describe(percentiles=percentiles).transpose()
@@ -1552,6 +1558,8 @@ async def main() -> None:
url=args.url, num_clients=args.num_clients, early_stop=not args.no_early_stop
)
+ warmup_runtime_sec: float | None = None
+
# Warm-up step
if args.warmup_step:
# Only send a single user prompt from every conversation.
@@ -1566,26 +1574,56 @@ async def main() -> None:
# all clients should finish their work before exiting
warmup_bench_args = bench_args._replace(early_stop=False)
- logger.info(f"{Color.PURPLE}Warmup start{Color.RESET}")
+ logger.info("%sWarmup start%s", Color.PURPLE, Color.RESET)
+ warmup_start_ns = time.perf_counter_ns()
conversations, _ = await main_mp(
warmup_client_args, req_args, warmup_bench_args, tokenizer, conversations
)
- logger.info(f"{Color.PURPLE}Warmup done{Color.RESET}")
+ warmup_runtime_sec = nanosec_to_sec(time.perf_counter_ns() - warmup_start_ns)
+ logger.info(
+ "%sWarmup runtime: %.3f sec (%.3f ms)%s",
+ Color.PURPLE,
+ warmup_runtime_sec,
+ warmup_runtime_sec * 1000,
+ Color.RESET,
+ )
+ logger.info("%sWarmup done%s", Color.PURPLE, Color.RESET)
# Run the benchmark
- start_time = time.perf_counter_ns()
+ benchmark_start_ns = time.perf_counter_ns()
client_convs, client_metrics = await main_mp(
client_args, req_args, bench_args, tokenizer, conversations
)
- total_runtime_ms = nanosec_to_millisec(time.perf_counter_ns() - start_time)
+ benchmark_runtime_sec = nanosec_to_sec(time.perf_counter_ns() - benchmark_start_ns)
# Calculate requests per second
- total_runtime_sec = total_runtime_ms / 1000.0
- rps = len(client_metrics) / total_runtime_sec
+ requests_per_sec = len(client_metrics) / benchmark_runtime_sec
+ benchmark_runtime_ms = benchmark_runtime_sec * 1000.0
logger.info(
- f"{Color.GREEN}All clients finished, total runtime: {total_runtime_sec:.3f} sec"
- f" ({total_runtime_ms:.3f} ms), requests per second: {rps:.3f}{Color.RESET}"
+ "%sAll clients finished, benchmark runtime: %.3f sec (%.3f ms), "
+ "requests per second: %.3f%s",
+ Color.GREEN,
+ benchmark_runtime_sec,
+ benchmark_runtime_ms,
+ requests_per_sec,
+ Color.RESET,
)
+ if warmup_runtime_sec is not None:
+ total_runtime_sec = benchmark_runtime_sec + warmup_runtime_sec
+ logger.info(
+ "%sWarmup runtime: %.3f sec (%.3f ms)%s",
+ Color.GREEN,
+ warmup_runtime_sec,
+ warmup_runtime_sec * 1000,
+ Color.RESET,
+ )
+ logger.info(
+ "%sTotal runtime (including warmup): %.3f sec (%.3f ms)%s",
+ Color.GREEN,
+ total_runtime_sec,
+ total_runtime_sec * 1000,
+ Color.RESET,
+ )
# Benchmark parameters
params = {
@@ -1610,6 +1648,7 @@ async def main() -> None:
verbose=args.verbose,
gen_conv_args=gen_conv_args,
excel_output=args.excel_output,
+ warmup_runtime_sec=warmup_runtime_sec,
)
if args.output_file is not None:
diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake
index aa84125818d10..fbbb03c5ed465 100644
--- a/cmake/cpu_extension.cmake
+++ b/cmake/cpu_extension.cmake
@@ -375,6 +375,7 @@ set(VLLM_EXT_SRC
if (AVX512_FOUND AND NOT AVX512_DISABLED)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
+ "csrc/cpu/cpu_wna16.cpp"
${VLLM_EXT_SRC})
if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
set(VLLM_EXT_SRC
diff --git a/cmake/external_projects/triton_kernels.cmake b/cmake/external_projects/triton_kernels.cmake
new file mode 100644
index 0000000000000..d35ad123dd9de
--- /dev/null
+++ b/cmake/external_projects/triton_kernels.cmake
@@ -0,0 +1,53 @@
+# Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
+
+set(DEFAULT_TRITON_KERNELS_TAG "v3.5.0")
+
+# Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
+# be directly set to the triton_kernels python directory.
+if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
+ message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
+ FetchContent_Declare(
+ triton_kernels
+ SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR}
+ )
+
+else()
+ set(TRITON_GIT "https://github.com/triton-lang/triton.git")
+ message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}")
+ FetchContent_Declare(
+ triton_kernels
+ # TODO (varun) : Fetch just the triton_kernels directory from Triton
+ GIT_REPOSITORY https://github.com/triton-lang/triton.git
+ GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG}
+ GIT_PROGRESS TRUE
+ SOURCE_SUBDIR python/triton_kernels/triton_kernels
+ )
+endif()
+
+# Fetch content
+FetchContent_MakeAvailable(triton_kernels)
+
+if (NOT triton_kernels_SOURCE_DIR)
+ message (FATAL_ERROR "[triton_kernels] Cannot resolve triton_kernels_SOURCE_DIR")
+endif()
+
+if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
+ set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/")
+else()
+ set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/python/triton_kernels/triton_kernels/")
+endif()
+
+message (STATUS "[triton_kernels] triton_kernels is available at ${TRITON_KERNELS_PYTHON_DIR}")
+
+add_custom_target(triton_kernels)
+
+# Ensure the vllm/third_party directory exists before installation
+install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/triton_kernels\")")
+
+## Copy .py files to install directory.
+install(DIRECTORY
+ ${TRITON_KERNELS_PYTHON_DIR}
+ DESTINATION
+ vllm/third_party/triton_kernels/
+ COMPONENT triton_kernels
+ FILES_MATCHING PATTERN "*.py")
diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake
index 567c8959f0454..ff687e0af7b44 100644
--- a/cmake/external_projects/vllm_flash_attn.cmake
+++ b/cmake/external_projects/vllm_flash_attn.cmake
@@ -38,7 +38,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
- GIT_TAG 58e0626a692f09241182582659e3bf8f16472659
+ GIT_TAG 86f8f157cf82aa2342743752b97788922dd7de43
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu
index 0aa0dc14c7480..32960cc8073bb 100644
--- a/csrc/cache_kernels.cu
+++ b/csrc/cache_kernels.cu
@@ -552,7 +552,11 @@ __global__ void indexer_k_quant_and_cache_kernel(
#ifndef USE_ROCM
__syncwarp();
#endif
+#if defined(__gfx942__)
+ float scale = fmaxf(amax, 1e-4) / 224.0f;
+#else
float scale = fmaxf(amax, 1e-4) / 448.0f;
+#endif
if (use_ue8m0) {
scale = exp2f(ceilf(log2f(scale)));
}
@@ -965,7 +969,9 @@ __global__ void gather_and_maybe_dequant_cache(
}
};
- for (int pid = split_start; pid < full_blocks_end; ++pid) {
+ const auto loop_end =
+ std::min((int64_t)full_blocks_end, block_table_stride - offset);
+ for (int pid = split_start; pid < loop_end; ++pid) {
auto block_id = batch_block_table[pid];
auto block_start_ptr = src_cache + block_id * cache_block_stride;
auto block_dst_ptr = dst + pid * block_size * dst_entry_stride;
@@ -976,12 +982,15 @@ __global__ void gather_and_maybe_dequant_cache(
}
if (partial_block_size) {
- auto block_id = batch_block_table[full_blocks_end];
- auto block_start_ptr = src_cache + block_id * cache_block_stride;
- auto block_dst_ptr = dst + full_blocks_end * block_size * dst_entry_stride;
- for (int eid = 0; eid < partial_block_size; ++eid) {
- copy_entry(block_start_ptr + eid * cache_entry_stride,
- block_dst_ptr + eid * dst_entry_stride);
+ if (offset + full_blocks_end < block_table_stride) {
+ auto block_id = batch_block_table[full_blocks_end];
+ auto block_start_ptr = src_cache + block_id * cache_block_stride;
+ auto block_dst_ptr =
+ dst + full_blocks_end * block_size * dst_entry_stride;
+ for (int eid = 0; eid < partial_block_size; ++eid) {
+ copy_entry(block_start_ptr + eid * cache_entry_stride,
+ block_dst_ptr + eid * dst_entry_stride);
+ }
}
}
}
diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp
index 50f17c758c148..92f8bee5a47a0 100644
--- a/csrc/cpu/cpu_attn.cpp
+++ b/csrc/cpu/cpu_attn.cpp
@@ -13,6 +13,18 @@
#define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
#endif
+#ifdef __aarch64__
+ #include "cpu_attn_neon.hpp"
+ #define NEON_DISPATCH(...) \
+ case cpu_attention::ISA::NEON: { \
+ using attn_impl = cpu_attention::AttentionImpl; \
+ return __VA_ARGS__(); \
+ }
+#else
+ #define NEON_DISPATCH(...) case cpu_attention::ISA::NEON:
+#endif // #ifdef __aarch64__
+
#define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \
case HEAD_DIM: { \
constexpr size_t head_dim = HEAD_DIM; \
@@ -41,6 +53,7 @@
[&] { \
switch (ISA_TYPE) { \
AMX_DISPATCH(__VA_ARGS__) \
+ NEON_DISPATCH(__VA_ARGS__) \
case cpu_attention::ISA::VEC: { \
using attn_impl = \
cpu_attention::AttentionImpl
#include
#include
@@ -12,9 +11,10 @@
#include "cpu_types.hpp"
#include "scratchpad_manager.h"
#include "cpu_attn_macros.h"
+#include "utils.hpp"
namespace cpu_attention {
-enum class ISA { AMX, VEC, VEC16 };
+enum class ISA { AMX, VEC, VEC16, NEON };
template
class AttentionImpl {};
@@ -143,6 +143,12 @@ struct AttentionMetadata {
case ISA::VEC:
ss << "VEC, ";
break;
+ case ISA::VEC16:
+ ss << "VEC16, ";
+ break;
+ case ISA::NEON:
+ ss << "NEON, ";
+ break;
}
ss << "workitem_group_num: " << workitem_group_num
<< ", reduction_item_num: " << reduction_item_num
diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp
new file mode 100644
index 0000000000000..827f0cfbc718e
--- /dev/null
+++ b/csrc/cpu/cpu_attn_neon.hpp
@@ -0,0 +1,386 @@
+#ifndef CPU_ATTN_NEON_HPP
+#define CPU_ATTN_NEON_HPP
+
+#include "cpu_attn_impl.hpp"
+#include
+#include
+namespace cpu_attention {
+
+namespace {
+
+#define BLOCK_SIZE_ALIGNMENT 32
+#define HEAD_SIZE_ALIGNMENT 32
+#define MAX_Q_HEAD_NUM_PER_ITER 16
+
+// These do not use vectorized class for loading / converting
+// because csrc/cpu/cpu_types_arm.hpp does not have fallback options
+// for vec_op::BF16Vec* / vec_op::BF16Vec* on Arm HW that
+// doesn't support BF16.
+// We don't use vec_op::FP32Vec* or vec_op::FP16Vec* for consistency.
+template
+FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, float32x4_t& b0,
+ float32x4_t& b1);
+
+template <>
+FORCE_INLINE void load_row8_B_as_f32(const float* p, float32x4_t& b0,
+ float32x4_t& b1) {
+ b0 = vld1q_f32(p + 0);
+ b1 = vld1q_f32(p + 4);
+}
+
+template <>
+FORCE_INLINE void load_row8_B_as_f32(const c10::Half* p,
+ float32x4_t& b0,
+ float32x4_t& b1) {
+ const float16_t* h = reinterpret_cast(p);
+ float16x8_t v = vld1q_f16(h);
+ b0 = vcvt_f32_f16(vget_low_f16(v));
+ b1 = vcvt_f32_f16(vget_high_f16(v));
+}
+
+template <>
+FORCE_INLINE void load_row8_B_as_f32(const c10::BFloat16* p,
+ float32x4_t& b0,
+ float32x4_t& b1) {
+ const uint16_t* u = reinterpret_cast(p);
+#ifdef ARM_BF16_SUPPORT
+ uint16x8_t u0 = vld1q_u16(u);
+ bfloat16x8_t bf0 = vreinterpretq_bf16_u16(u0);
+ b0 = vcvtq_low_f32_bf16(bf0);
+ b1 = vcvtq_high_f32_bf16(bf0);
+#else
+ uint16x8_t x0 = vld1q_u16(u);
+ uint32x4_t lo = vshlq_n_u32(vmovl_u16(vget_low_u16(x0)), 16);
+ uint32x4_t hi = vshlq_n_u32(vmovl_u16(vget_high_u16(x0)), 16);
+ b0 = vreinterpretq_f32_u32(lo);
+ b1 = vreinterpretq_f32_u32(hi);
+#endif
+}
+
+// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with NEON FMLAs
+// #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2)
+// #FMLAs = (K // 4) * (4 * 2 * M)
+// We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads
+template
+FORCE_INLINE void gemm_micro_neon_fmla_Mx8_Ku4(
+ const float* __restrict A, // [M x K],
+ const kv_cache_t* __restrict B, // [K x 8],
+ float* __restrict C, // [M x 8],
+ int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
+ // kernel supports max M of 8, as it'd spill for larger M
+ static_assert(1 <= M && M <= 8, "M must be in [1,8]");
+
+// helpers for per-M codegen
+#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
+#define IF_M(i) if constexpr (M > (i))
+
+ // A row base pointers
+#define DECL_A(i) const float* a##i = A + (i) * lda;
+ ROWS_APPLY(DECL_A)
+#undef DECL_A
+
+ // declare 2 accumulators per row of M
+#define DECL_ACC(i) float32x4_t acc##i##_0, acc##i##_1;
+ ROWS_APPLY(DECL_ACC)
+#undef DECL_ACC
+
+ // initialize accumulators
+#define INIT_ACC(i) \
+ IF_M(i) { \
+ if (accumulate) { \
+ acc##i##_0 = vld1q_f32(C + (i) * ldc + 0); \
+ acc##i##_1 = vld1q_f32(C + (i) * ldc + 4); \
+ } else { \
+ acc##i##_0 = vdupq_n_f32(0.f); \
+ acc##i##_1 = vdupq_n_f32(0.f); \
+ } \
+ }
+ ROWS_APPLY(INIT_ACC)
+#undef INIT_ACC
+
+ int32_t k = 0;
+
+ // K unrolled by 4
+ for (; k + 3 < K; k += 4) {
+ // load A[k..k+3] for each active row (M)
+#define LOAD_A4(i) \
+ float32x4_t a##i##v; \
+ IF_M(i) a##i##v = vld1q_f32(a##i + k);
+ ROWS_APPLY(LOAD_A4)
+#undef LOAD_A4
+
+ // helper: FMA lane L from aiv
+#define FMAS_LANE(i, aiv, L) \
+ IF_M(i) { \
+ acc##i##_0 = vfmaq_laneq_f32(acc##i##_0, b0, aiv, L); \
+ acc##i##_1 = vfmaq_laneq_f32(acc##i##_1, b1, aiv, L); \
+ }
+
+ // k + 0
+ {
+ float32x4_t b0, b1;
+ load_row8_B_as_f32(B + (int64_t)(k + 0) * ldb, b0, b1);
+#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
+ ROWS_APPLY(STEP_K0)
+#undef STEP_K0
+ }
+ // k + 1
+ {
+ float32x4_t b0, b1;
+ load_row8_B_as_f32(B + (int64_t)(k + 1) * ldb, b0, b1);
+#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
+ ROWS_APPLY(STEP_K1)
+#undef STEP_K1
+ }
+ // k + 2
+ {
+ float32x4_t b0, b1;
+ load_row8_B_as_f32(B + (int64_t)(k + 2) * ldb, b0, b1);
+#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
+ ROWS_APPLY(STEP_K2)
+#undef STEP_K2
+ }
+ // k + 3
+ {
+ float32x4_t b0, b1;
+ load_row8_B_as_f32(B + (int64_t)(k + 3) * ldb, b0, b1);
+#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
+ ROWS_APPLY(STEP_K3)
+#undef STEP_K3
+ }
+#undef FMAS_LANE
+ }
+
+ // K tail
+ for (; k < K; ++k) {
+ float32x4_t b0, b1;
+ load_row8_B_as_f32(B + (int64_t)k * ldb, b0, b1);
+#define TAIL_ROW(i) \
+ IF_M(i) { \
+ float32x4_t ai = vdupq_n_f32(*(a##i + k)); \
+ acc##i##_0 = vfmaq_f32(acc##i##_0, b0, ai); \
+ acc##i##_1 = vfmaq_f32(acc##i##_1, b1, ai); \
+ }
+ ROWS_APPLY(TAIL_ROW)
+#undef TAIL_ROW
+ }
+
+ // store accumulators to C
+#define STORE_ROW(i) \
+ IF_M(i) { \
+ vst1q_f32(C + (i) * ldc + 0, acc##i##_0); \
+ vst1q_f32(C + (i) * ldc + 4, acc##i##_1); \
+ }
+ ROWS_APPLY(STORE_ROW)
+#undef STORE_ROW
+
+#undef ROWS_APPLY
+#undef IF_M
+}
+
+template
+FORCE_INLINE void gemm_macro_neon_fmla_Mx8_Ku4(const float* __restrict A,
+ const kv_cache_t* __restrict B,
+ float* __restrict C, int32_t M,
+ int32_t K, int64_t lda,
+ int64_t ldb, int64_t ldc,
+ bool accumulate) {
+ // micro kernel is Mx8
+ static_assert(N % 8 == 0, "N must be a multiple of 8");
+ for (int32_t m = 0; m < M;) {
+ int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
+ const float* Ab = A + m * lda;
+ float* Cb = C + m * ldc;
+
+ for (int32_t n = 0; n < N; n += 8) {
+ const kv_cache_t* Bn = B + n;
+ float* Cn = Cb + n;
+ switch (mb) {
+ case 8:
+ gemm_micro_neon_fmla_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+ K, accumulate);
+ break;
+ case 4:
+ gemm_micro_neon_fmla_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+ K, accumulate);
+ break;
+ case 2:
+ gemm_micro_neon_fmla_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+ K, accumulate);
+ break;
+ default:
+ gemm_micro_neon_fmla_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+ K, accumulate);
+ break;
+ }
+ }
+ // no tail loop for N as it's guaranteed to be a multiple of 8
+ m += mb;
+ }
+}
+
+template
+class TileGemmNeonFMLA {
+ public:
+ template
+ FORCE_INLINE static void gemm(const int32_t m_size,
+ float* __restrict__ a_tile,
+ kv_cache_t* __restrict__ b_tile,
+ float* __restrict__ c_tile, const int64_t lda,
+ const int64_t ldb, const int64_t ldc,
+ const int32_t block_size,
+ const int32_t dynamic_k_size,
+ const bool accum_c) {
+ if constexpr (phase == AttentionGemmPhase::QK) {
+ gemm_macro_neon_fmla_Mx8_Ku4(
+ a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
+ } else {
+ gemm_macro_neon_fmla_Mx8_Ku4(
+ a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
+ accum_c);
+ }
+ }
+};
+
+} // namespace
+
+// this is similar to "ISA::VEC" at the moment
+template
+class AttentionImpl {
+ public:
+ using query_t = scalar_t;
+ using q_buffer_t = float;
+ using kv_cache_t = scalar_t;
+ using logits_buffer_t = float;
+ using partial_output_buffer_t = float;
+ using prob_buffer_t = float;
+
+ constexpr static int64_t BlockSizeAlignment =
+ BLOCK_SIZE_ALIGNMENT; // KV token num unit of QK and PV phases
+ constexpr static int64_t HeadDimAlignment =
+ HEAD_SIZE_ALIGNMENT; // headdim num unit of PV phase
+ constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
+ constexpr static int64_t HeadDim = head_dim;
+ constexpr static ISA ISAType = ISA::NEON;
+ constexpr static bool scale_on_logits = false; // apply scale on q_buffer
+
+ static_assert(HeadDim % HeadDimAlignment == 0);
+ // the gemm micro kernel is Mx8
+ static_assert(HeadDimAlignment % 8 == 0);
+ static_assert(BlockSizeAlignment % 8 == 0);
+
+ public:
+ template typename attention>
+ FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
+ attention> attention_iteration;
+ attention_iteration(CPU_ATTENTION_PARAMS);
+ }
+
+ // k_cache_token_group_stride: stride of K cache when move to next
+ // BlockSizeAlignment tokens in a block
+ constexpr static int64_t k_cache_token_group_stride(
+ const int32_t block_size) {
+ return BlockSizeAlignment; // layout of k_cache block is [head_dim,
+ // block_size], row-major
+ }
+
+ // v_cache_token_group_stride: stride of V cache when move to next
+ // BlockSizeAlignment tokens in a block
+ constexpr static int64_t v_cache_token_group_stride(
+ const int32_t block_size) {
+ return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
+ // head_dim], row-major
+ }
+
+ // v_cache_head_group_stride: stride of V cache when move to next
+ // HeadDimAlignment head dims in a block
+ constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
+ return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
+ // row-major
+ }
+
+ // Copy q to q_buffer and cast it to fp32
+ static void copy_q_heads_tile(
+ scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
+ float* __restrict__ q_buffer, const int32_t q_num,
+ const int32_t q_heads_per_kv, const int64_t q_num_stride,
+ const int64_t q_head_stride, float scale) {
+ static_assert(head_dim % 16 == 0);
+ constexpr int32_t unroll_size = head_dim / 16;
+ using load_vec_t = typename VecTypeTrait::vec_t;
+
+ vec_op::FP32Vec16 scale_vec(scale);
+ for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
+ for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
+ scalar_t* __restrict__ curr_q =
+ src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
+ float* __restrict__ curr_q_buffer =
+ q_buffer + q_num_idx * q_heads_per_kv * head_dim +
+ q_head_idx * head_dim;
+
+ vec_op::unroll_loop([&](int32_t i) {
+ load_vec_t vec(curr_q);
+ vec_op::FP32Vec16 fp32_vec(vec);
+ fp32_vec = fp32_vec * scale_vec;
+ fp32_vec.save(curr_q_buffer);
+
+ curr_q += 16;
+ curr_q_buffer += 16;
+ });
+ }
+ }
+ }
+
+ // reshape K as column-major and V as row-major
+ static void reshape_and_cache(
+ const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
+ scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
+ const int64_t* __restrict__ slot_mapping, const int64_t token_num,
+ const int64_t key_token_num_stride, const int64_t value_token_num_stride,
+ const int64_t head_num, const int64_t key_head_num_stride,
+ const int64_t value_head_num_stride, const int64_t num_blocks,
+ const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
+ const int64_t block_size, const int64_t block_size_stride) {
+#pragma omp parallel for collapse(2)
+ for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+ for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
+ const int64_t pos = slot_mapping[token_idx];
+ if (pos < 0) {
+ // skip
+ continue;
+ }
+
+ const int64_t block_idx = pos / block_size;
+ const int64_t block_offset = pos % block_size;
+ {
+ // Write Key
+ const scalar_t* key_start_ptr = key +
+ token_idx * key_token_num_stride +
+ head_idx * key_head_num_stride;
+ scalar_t* key_cache_start_ptr =
+ key_cache + block_idx * num_blocks_stride +
+ head_idx * cache_head_num_stride + block_offset;
+
+#pragma GCC unroll 8
+ for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
+ key_cache_start_ptr[j] = key_start_ptr[i];
+ }
+ }
+ {
+ // Write Value
+ const scalar_t* value_start_ptr = value +
+ token_idx * value_token_num_stride +
+ head_idx * value_head_num_stride;
+ scalar_t* value_cache_start_ptr =
+ value_cache + block_idx * num_blocks_stride +
+ head_idx * cache_head_num_stride + block_offset * head_dim;
+ std::memcpy(value_cache_start_ptr, value_start_ptr,
+ sizeof(scalar_t) * head_dim);
+ }
+ }
+ }
+ }
+};
+} // namespace cpu_attention
+
+#endif // #ifndef CPU_ATTN_NEON_HPP
diff --git a/csrc/cpu/cpu_types_scalar.hpp b/csrc/cpu/cpu_types_scalar.hpp
index 1a9278bc662e5..f9da78283da5e 100644
--- a/csrc/cpu/cpu_types_scalar.hpp
+++ b/csrc/cpu/cpu_types_scalar.hpp
@@ -26,10 +26,6 @@ namespace vec_op {
#define FORCE_INLINE __attribute__((always_inline)) inline
-#define __max(a, b) ((a) > (b) ? (a) : (b))
-#define __min(a, b) ((a) < (b) ? (a) : (b))
-#define __abs(a) ((a) < (0) ? (0 - a) : (a))
-
typedef struct f16x8_t {
uint16_t val[8];
} f16x8_t;
@@ -99,7 +95,7 @@ struct FP16Vec16 : public Vec {
void save(void* ptr) const { *reinterpret_cast(ptr) = reg; }
void save(void* ptr, const int elem_num) const {
- int num = __min(elem_num, VEC_ELEM_NUM);
+ int num = std::min(elem_num, VEC_ELEM_NUM);
std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
}
};
@@ -128,7 +124,7 @@ struct BF16Vec16 : public Vec {
void save(void* ptr) const { *reinterpret_cast(ptr) = reg; }
void save(void* ptr, const int elem_num) const {
- int num = __min(elem_num, VEC_ELEM_NUM);
+ int num = std::min(elem_num, VEC_ELEM_NUM);
std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
}
};
@@ -143,9 +139,9 @@ struct BF16Vec32 : public Vec {
explicit BF16Vec32(f16x32_t data) : reg(data) {};
explicit BF16Vec32(BF16Vec8& vec8_data) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+ unroll_loop([&vec8_data, this](int i) {
reg.val[i] = vec8_data.reg.val[i % BF16Vec8::VEC_ELEM_NUM];
- }
+ });
}
void save(void* ptr) const { *reinterpret_cast(ptr) = reg; }
@@ -157,15 +153,11 @@ struct FP32Vec4 : public Vec {
f32x4_t reg;
explicit FP32Vec4(float v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = v;
- }
+ unroll_loop([&v, this](int i) { reg.val[i] = v; });
}
explicit FP32Vec4() {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = 0.0f;
- }
+ unroll_loop([this](int i) { reg.val[i] = 0.0f; });
}
explicit FP32Vec4(const float* ptr)
@@ -182,15 +174,11 @@ struct FP32Vec8 : public Vec {
f32x8_t reg;
explicit FP32Vec8(float v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = v;
- }
+ unroll_loop([&v, this](int i) { reg.val[i] = v; });
}
explicit FP32Vec8() {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = 0.0f;
- }
+ unroll_loop([this](int i) { reg.val[i] = 0.0f; });
}
explicit FP32Vec8(const float* ptr)
@@ -201,78 +189,68 @@ struct FP32Vec8 : public Vec {
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
explicit FP32Vec8(const FP16Vec8& v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = fp16_to_float(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
}
FP32Vec8(const BF16Vec8& v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = bf16_to_float(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
}
float reduce_sum() const {
float result = 0;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result += reg.val[i];
- }
+ unroll_loop(
+ [&result, this](int i) { result += reg.val[i]; });
return result;
}
FP32Vec8 exp() const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = expf(reg.val[i]);
- }
+ unroll_loop(
+ [&ret, this](int i) { ret.val[i] = expf(reg.val[i]); });
return FP32Vec8(ret);
}
FP32Vec8 tanh() const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = tanhf(reg.val[i]);
- }
+ unroll_loop(
+ [&ret, this](int i) { ret.val[i] = tanhf(reg.val[i]); });
return FP32Vec8(ret);
}
FP32Vec8 er() const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = erf(reg.val[i]);
- }
+ unroll_loop(
+ [&ret, this](int i) { ret.val[i] = erf(reg.val[i]); });
return FP32Vec8(ret);
}
FP32Vec8 operator*(const FP32Vec8& b) const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = reg.val[i] * b.reg.val[i];
- }
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
return FP32Vec8(ret);
}
FP32Vec8 operator+(const FP32Vec8& b) const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = reg.val[i] + b.reg.val[i];
- }
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
return FP32Vec8(ret);
}
FP32Vec8 operator-(const FP32Vec8& b) const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = reg.val[i] - b.reg.val[i];
- }
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
return FP32Vec8(ret);
}
FP32Vec8 operator/(const FP32Vec8& b) const {
f32x8_t ret;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- ret.val[i] = reg.val[i] / b.reg.val[i];
- }
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
return FP32Vec8(ret);
}
@@ -284,15 +262,11 @@ struct FP32Vec16 : public Vec {
f32x16_t reg;
explicit FP32Vec16(float v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = v;
- }
+ unroll_loop([&v, this](int i) { reg.val[i] = v; });
}
explicit FP32Vec16() {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = 0.0f;
- }
+ unroll_loop([this](int i) { reg.val[i] = 0.0f; });
}
explicit FP32Vec16(const float* ptr)
@@ -301,29 +275,27 @@ struct FP32Vec16 : public Vec {
explicit FP32Vec16(f32x16_t data) : reg(data) {};
FP32Vec16(const FP32Vec4& data) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+ unroll_loop([&data, this](int i) {
reg.val[i] = data.reg.val[i % FP32Vec4::VEC_ELEM_NUM];
- }
+ });
}
FP32Vec16(const FP32Vec8& data) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+ unroll_loop([&data, this](int i) {
reg.val[i] = data.reg.val[i % FP32Vec8::VEC_ELEM_NUM];
- }
+ });
}
FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(const FP16Vec16& v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = fp16_to_float(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
}
explicit FP32Vec16(const BF16Vec16& v) {
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- reg.val[i] = bf16_to_float(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
}
explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
@@ -331,82 +303,74 @@ struct FP32Vec16 : public Vec {
FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
FP32Vec16 operator*(const FP32Vec16& b) const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = reg.val[i] * b.reg.val[i];
- }
- return result;
+ f32x16_t ret;
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
+ return FP32Vec16(ret);
}
FP32Vec16 operator+(const FP32Vec16& b) const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = reg.val[i] + b.reg.val[i];
- }
- return result;
+ f32x16_t ret;
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
+ return FP32Vec16(ret);
}
FP32Vec16 operator-(const FP32Vec16& b) const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = reg.val[i] - b.reg.val[i];
- }
- return result;
+ f32x16_t ret;
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
+ return FP32Vec16(ret);
}
FP32Vec16 operator/(const FP32Vec16& b) const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = reg.val[i] / b.reg.val[i];
- }
- return result;
+ f32x16_t ret;
+ unroll_loop(
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
+ return FP32Vec16(ret);
}
FP32Vec16 max(const FP32Vec16& b) const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = __max(reg.val[i], b.reg.val[i]);
- }
- return result;
+ f32x16_t ret;
+ unroll_loop([&ret, &b, this](int i) {
+ ret.val[i] = std::max(reg.val[i], b.reg.val[i]);
+ });
+ return FP32Vec16(ret);
}
FP32Vec16 min(const FP32Vec16& b) const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = __min(reg.val[i], b.reg.val[i]);
- }
- return result;
+ f32x16_t ret;
+ unroll_loop([&ret, &b, this](int i) {
+ ret.val[i] = std::min(reg.val[i], b.reg.val[i]);
+ });
+ return FP32Vec16(ret);
}
FP32Vec16 abs() const {
- FP32Vec16 result(0.0f);
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result.reg.val[i] = __abs(reg.val[i]);
- }
- return result;
+ f32x16_t ret;
+ unroll_loop(
+ [&ret, this](int i) { ret.val[i] = std::abs(reg.val[i]); });
+ return FP32Vec16(ret);
}
float reduce_sum() const {
float result = 0.0f;
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result += reg.val[i];
- }
+ unroll_loop(
+ [&result, this](int i) { result += reg.val[i]; });
return result;
}
float reduce_max() const {
- float result = reg.val[0];
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result = __max(reg.val[i], result);
- }
+ float result = std::numeric_limits::lowest();
+ unroll_loop(
+ [&result, this](int i) { result = std::max(reg.val[i], result); });
return result;
}
float reduce_min() const {
- float result = reg.val[0];
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
- result = __min(reg.val[i], result);
- }
+ float result = std::numeric_limits::max();
+ unroll_loop(
+ [&result, this](int i) { result = std::min(reg.val[i], result); });
return result;
}
@@ -414,13 +378,9 @@ struct FP32Vec16 : public Vec {
float reduce_sub_sum(int idx) {
static_assert(VEC_ELEM_NUM % group_size == 0);
float sum = 0.0;
- int start = idx * group_size;
- int end = (idx + 1) * group_size;
-
- for (; (start < VEC_ELEM_NUM) && (start < end); ++start) {
- sum += reg.val[start];
- }
-
+ const int start = idx * group_size;
+ unroll_loop(
+ [&sum, &start, this](int i) { sum += reg.val[start + i]; });
return sum;
}
@@ -477,17 +437,13 @@ inline void storeFP32(float v, c10::BFloat16* ptr) {
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
- int i = 0;
- for (i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) {
- reg.val[i] = float_to_fp16(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
}
inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
- int i = 0;
- for (i = 0; i < FP16Vec8::VEC_ELEM_NUM; ++i) {
- reg.val[i] = float_to_fp16(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
}
inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
@@ -495,17 +451,13 @@ inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
}
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
- int i = 0;
- for (i = 0; i < BF16Vec8::VEC_ELEM_NUM; ++i) {
- reg.val[i] = float_to_bf16(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
- int i = 0;
- for (i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) {
- reg.val[i] = float_to_bf16(v.reg.val[i]);
- }
+ unroll_loop(
+ [&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
}
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 3); }
diff --git a/csrc/cpu/cpu_types_x86.hpp b/csrc/cpu/cpu_types_x86.hpp
index 7ddf028e6e131..6f51277f78440 100644
--- a/csrc/cpu/cpu_types_x86.hpp
+++ b/csrc/cpu/cpu_types_x86.hpp
@@ -104,6 +104,8 @@ struct FP16Vec16 : public Vec {
explicit FP16Vec16(bool, void* ptr)
: reg(_mm256_stream_load_si256((__m256i*)ptr)) {}
+ explicit FP16Vec16(const c10::Half v) : reg(_mm256_set1_epi16(v.x)) {}
+
explicit FP16Vec16(const FP32Vec16&);
void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
@@ -141,6 +143,8 @@ struct BF16Vec16 : public Vec {
explicit BF16Vec16(bool, void* ptr)
: reg(_mm256_stream_load_si256((__m256i*)ptr)) {}
+ explicit BF16Vec16(const c10::BFloat16 v) : reg(_mm256_set1_epi16(v.x)) {}
+
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
@@ -350,6 +354,22 @@ struct FP32Vec16 : public Vec {
explicit FP32Vec16(__m512 data) : reg(data) {}
+ // de-pack 4 bit values
+ explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
+ int64_t mask_0 = 0x0F0F0F0F0F0F0F0F;
+ int64_t mask_1 = 0xF0F0F0F0F0F0F0F0;
+ int64_t value_0 = value & mask_0;
+ int64_t value_1 = value & mask_1;
+ __m128i vec_0 = _mm_movpi64_epi64((__m64)value_0);
+ __m128i vec_1 = _mm_movpi64_epi64((__m64)value_1);
+ vec_0 = _mm_cvtepu8_epi16(vec_0);
+ vec_1 = _mm_cvtepu8_epi16(vec_1);
+ vec_1 = _mm_slli_epi16(vec_1, 4);
+ __m128i vec = _mm_or_si128(vec_0, vec_1);
+ __m512i vec_i32 = _mm512_cvtepu8_epi32(vec);
+ reg = _mm512_permutexvar_ps(vec_i32, lut.reg);
+ }
+
explicit FP32Vec16(const FP32Vec4& data)
: reg((__m512)_mm512_inserti32x4(
_mm512_inserti32x4(
@@ -426,14 +446,6 @@ struct FP32Vec16 : public Vec {
float get_last_elem() const { return _mm512_cvtss_f32(reg); }
- template
- float reduce_sub_sum(int idx) {
- static_assert(VEC_ELEM_NUM % group_size == 0);
- constexpr uint32_t base_mask = (0xFFFF >> (16 - group_size));
- __mmask16 mask = _cvtu32_mask16(base_mask << (idx * group_size));
- return _mm512_mask_reduce_add_ps(mask, reg);
- }
-
void save(float* ptr) const { _mm512_storeu_ps(ptr, reg); }
void save(float* ptr, const int elem_num) const {
@@ -755,6 +767,25 @@ inline void non_temporal_save(BF16Vec16& vec, void* ptr) {
inline void non_temporal_save(FP32Vec16& vec, void* ptr) {
_mm512_stream_ps((float*)ptr, vec.reg);
}
+
+static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
+ void* ptr) {
+ __m512i vec_0 = _mm512_cvtepu16_epi32(vec0.reg);
+ __m512i vec_1 = _mm512_cvtepu16_epi32(vec1.reg);
+ vec_1 = _mm512_slli_epi32(vec_1, 16);
+ vec_0 = _mm512_or_si512(vec_0, vec_1);
+ _mm512_storeu_epi32(ptr, vec_0);
+}
+
+static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
+ void* ptr) {
+ __m512i vec_0 = _mm512_cvtepu16_epi32(vec0.reg);
+ __m512i vec_1 = _mm512_cvtepu16_epi32(vec1.reg);
+ vec_1 = _mm512_slli_epi32(vec_1, 16);
+ vec_0 = _mm512_or_si512(vec_0, vec_1);
+ _mm512_storeu_epi32(ptr, vec_0);
+}
+
#endif
inline void mem_barrier() { _mm_mfence(); }
diff --git a/csrc/cpu/cpu_wna16.cpp b/csrc/cpu/cpu_wna16.cpp
new file mode 100644
index 0000000000000..816d195506e52
--- /dev/null
+++ b/csrc/cpu/cpu_wna16.cpp
@@ -0,0 +1,402 @@
+#include "cpu_types.hpp"
+#include "scratchpad_manager.h"
+#include "utils.hpp"
+
+#ifdef CPU_CAPABILITY_AMXBF16
+ #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
+#endif
+#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
+
+#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
+ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
+
+#define VLLM_DISPATCH_16B_TYPES(TYPE, NAME, ...) \
+ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_16B_TYPES(__VA_ARGS__))
+
+template
+void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
+ int32_t stride) {
+ std::stringstream ss;
+ ss << std::fixed << std::setprecision(5) << name << ": [\n";
+ auto* curr_logits_buffer = ptr;
+ for (int32_t m = 0; m < row; ++m) {
+ for (int32_t n = 0; n < col; ++n) {
+ ss << curr_logits_buffer[n] << ", ";
+ }
+ ss << "\n";
+ curr_logits_buffer += stride;
+ }
+ ss << "]\n";
+ std::printf("%s", ss.str().c_str());
+}
+
+namespace {
+using cpu_utils::ISA;
+using cpu_utils::VecTypeTrait;
+
+template
+class Dequantizer4b {
+ public:
+ constexpr static int32_t pack_num = 32 / 4;
+ using scalar_vec_t = typename VecTypeTrait::vec_t;
+
+ public:
+ static void dequant(int32_t* __restrict__ q_weight,
+ scalar_t* __restrict__ weight,
+ scalar_t* __restrict__ scales,
+ int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
+ const int64_t scales_stride, const int64_t zeros_stride,
+ const int32_t k_size, const int32_t group_size) {
+ vec_op::FP32Vec16 lut;
+ if constexpr (has_zp) {
+ // AWQ
+ alignas(64) static const float LUT[16] = {
+ 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f,
+ 8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f};
+ lut = vec_op::FP32Vec16(LUT);
+ } else {
+ // GPTQ
+ alignas(64) static const float LUT[16] = {
+ -8.0f, -7.0f, -6.0f, -5.0f, -4.0f, -3.0f, -2.0f, -1.0f,
+ 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+ lut = vec_op::FP32Vec16(LUT);
+ }
+
+ // per 64-bits elem contains 16 output channels
+ int64_t* __restrict__ curr_q_weight = reinterpret_cast(q_weight);
+ int64_t* __restrict__ curr_zeros = reinterpret_cast(zeros);
+ scalar_t* __restrict__ curr_weight = weight;
+ scalar_t* __restrict__ curr_scale = scales;
+ vec_op::FP32Vec16 scale_0;
+ vec_op::FP32Vec16 scale_1;
+ vec_op::FP32Vec16 zero_0;
+ vec_op::FP32Vec16 zero_1;
+ int32_t group_counter = 0;
+ for (int32_t k_idx = 0; k_idx < k_size; k_idx += 2) {
+ int64_t qwb_0 = *curr_q_weight;
+ int64_t qwb_1 = *(curr_q_weight + 1);
+ vec_op::FP32Vec16 wb_0(qwb_0, lut);
+ vec_op::FP32Vec16 wb_1(qwb_1, lut);
+
+ if constexpr (!use_desc_act) {
+ if (group_counter == 0) {
+ scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale));
+ scale_1 = vec_op::FP32Vec16(scale_0);
+ curr_scale += scales_stride;
+
+ if constexpr (has_zp) {
+ zero_0 = vec_op::FP32Vec16(*curr_zeros, lut);
+ zero_1 = vec_op::FP32Vec16(zero_0);
+ curr_zeros += zeros_stride / 2;
+ }
+ }
+ } else {
+ int32_t g_idx_0 = g_idx[k_idx];
+ int32_t g_idx_1 = g_idx[k_idx + 1];
+ scale_0 = vec_op::FP32Vec16(
+ scalar_vec_t(curr_scale + g_idx_0 * scales_stride));
+ scale_1 = vec_op::FP32Vec16(
+ scalar_vec_t(curr_scale + g_idx_1 * scales_stride));
+ if constexpr (has_zp) {
+ zero_0 = vec_op::FP32Vec16(*(curr_zeros + g_idx_0 * zeros_stride / 2),
+ lut);
+ zero_1 = vec_op::FP32Vec16(*(curr_zeros + g_idx_1 * zeros_stride / 2),
+ lut);
+ }
+ }
+
+ if constexpr (has_zp) {
+ wb_0 = wb_0 - zero_0;
+ wb_1 = wb_1 - zero_1;
+ }
+
+ wb_0 = wb_0 * scale_0;
+ wb_1 = wb_1 * scale_1;
+
+ scalar_vec_t output_vec_0(wb_0);
+ scalar_vec_t output_vec_1(wb_1);
+
+ // AMX needs to interlave K elements to pack as 32 bits
+ if constexpr (isa == ISA::AMX) {
+ vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight);
+ } else {
+ output_vec_0.save(curr_weight);
+ output_vec_1.save(curr_weight + 16);
+ }
+
+ // update
+ curr_q_weight += 2;
+ curr_weight += 32;
+ if constexpr (!use_desc_act) {
+ group_counter += 2;
+ if (group_counter == group_size) {
+ group_counter = 0;
+ }
+ }
+ }
+ }
+};
+}; // namespace
+
+template
+void cpu_gemm_wna16_impl(
+ scalar_t* __restrict__ input, int32_t* __restrict__ q_weight,
+ scalar_t* __restrict__ output, scalar_t* __restrict__ scales,
+ int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
+ scalar_t* __restrict__ bias, const int32_t m_size, const int32_t n_size,
+ const int32_t k_size, const int64_t input_stride,
+ const int64_t output_stride, const int64_t scales_group_stride,
+ const int64_t zeros_group_stride, const int32_t group_num,
+ const int32_t group_size, const int64_t pack_factor) {
+ constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
+ 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();
+
+ // a simple schedule policy, just to hold more B tiles in L2 and make sure
+ // each thread has tasks
+ const int32_t n_partition_size = [&]() {
+ const int64_t cache_size = cpu_utils::get_l2_size();
+ int64_t ps_cache_limit = cache_size / (k_size * sizeof(scalar_t));
+ int64_t ps_thread_limit = n_size / thread_num;
+ ps_cache_limit =
+ std::max((ps_cache_limit / gemm_n_tile_size) * gemm_n_tile_size,
+ (int64_t)gemm_n_tile_size);
+ ps_thread_limit =
+ std::max((ps_thread_limit / gemm_n_tile_size) * gemm_n_tile_size,
+ (int64_t)gemm_n_tile_size);
+ return std::min(ps_cache_limit, ps_thread_limit);
+ }();
+ const int32_t task_num = (n_size + n_partition_size - 1) / n_partition_size;
+
+ // get buffer size
+ const int64_t b_buffer_size =
+ (((n_partition_size * k_size * sizeof(scalar_t) + 63) / 64) * 64);
+ const int64_t c_buffer_size =
+ (((gemm_m_tile_size * gemm_n_tile_size * sizeof(float) + 63) / 64) * 64);
+ const int64_t b_buffer_offset = 0;
+ const int64_t c_buffer_offset = b_buffer_size;
+ const int64_t buffer_size = b_buffer_size + c_buffer_size;
+ DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(buffer_size *
+ thread_num);
+
+ alignas(64) cpu_utils::Counter counter;
+ cpu_utils::Counter* counter_ptr = &counter;
+
+#pragma omp parallel for schedule(static, 1)
+ for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
+ scalar_t* __restrict__ b_buffer = nullptr;
+ float* __restrict__ c_buffer = nullptr;
+ {
+ uint8_t* buffer_ptr = DNNLScratchPadManager::get_dnnl_scratchpad_manager()
+ ->get_data() +
+ thread_id * buffer_size;
+ b_buffer = reinterpret_cast(buffer_ptr + b_buffer_offset);
+ c_buffer = reinterpret_cast(buffer_ptr + c_buffer_offset);
+ }
+
+ const int64_t q_weight_block_stride = n_block_size / pack_factor * k_size;
+ const int64_t b_buffer_block_stride = n_block_size * k_size;
+ const int32_t zeros_block_stride = n_block_size / pack_factor;
+
+ gemm_t gemm;
+
+ for (;;) {
+ int32_t task_id = counter_ptr->acquire_counter();
+
+ if (task_id >= task_num) {
+ break;
+ }
+
+ const int32_t n_start_idx = task_id * n_partition_size;
+ const int32_t n_block_start_idx = n_start_idx / n_block_size;
+ const int32_t n_num = std::min(n_partition_size, n_size - n_start_idx);
+ const int32_t n_block_num = n_num / n_block_size;
+ // std::printf("thread_id: %d, task_id: %d, n_start_idx: %d, n_num: %d\n",
+ // thread_id, task_id, n_start_idx, n_num);
+
+ // dequant weight
+ {
+ int32_t* __restrict__ curr_q_weight =
+ q_weight + n_block_start_idx * q_weight_block_stride;
+ scalar_t* __restrict__ curr_b_buffer = b_buffer;
+ scalar_t* __restrict__ curr_scales = scales + n_start_idx;
+ int32_t* __restrict__ curr_zeros = zeros + n_start_idx / pack_factor;
+ for (int32_t block_idx = 0; block_idx < n_block_num; ++block_idx) {
+ dequantizer_t::dequant(curr_q_weight, curr_b_buffer, curr_scales,
+ curr_zeros, g_idx, scales_group_stride,
+ zeros_group_stride, k_size, group_size);
+
+ // if (block_idx == 0 && n_start_idx == 0) {
+ // print_logits("depacked weight", curr_b_buffer, k_size,
+ // n_block_size, n_block_size);
+ // }
+
+ // update
+ curr_q_weight += q_weight_block_stride;
+ curr_b_buffer += b_buffer_block_stride;
+ curr_scales += n_block_size;
+ curr_zeros += zeros_block_stride;
+ }
+ }
+
+ // compute loop
+ {
+ const int32_t n_tile_num = n_num / gemm_n_tile_size;
+ scalar_t* __restrict__ curr_input = input;
+ scalar_t* __restrict__ init_bias = bias;
+ if (bias != nullptr) {
+ init_bias += n_start_idx;
+ }
+ scalar_t* __restrict__ init_output = output + n_start_idx;
+ for (int32_t m_idx = 0; m_idx < m_size; m_idx += gemm_m_tile_size) {
+ const int32_t curr_m_size =
+ std::min(gemm_m_tile_size, m_size - m_idx);
+ scalar_t* __restrict__ curr_b_buffer = b_buffer;
+ scalar_t* __restrict__ curr_bias = init_bias;
+ scalar_t* __restrict__ curr_output = init_output;
+ for (int32_t n_tile_idx = 0; n_tile_idx < n_tile_num; ++n_tile_idx) {
+ gemm.gemm(curr_input, curr_b_buffer, c_buffer, curr_m_size, k_size,
+ input_stride, b_buffer_block_stride, gemm_n_tile_size,
+ false);
+
+ if (bias != nullptr) {
+ cpu_micro_gemm::bias_epilogue(
+ c_buffer, curr_output, curr_bias, curr_m_size,
+ gemm_n_tile_size, output_stride);
+ curr_bias += gemm_n_tile_size;
+ } else {
+ cpu_micro_gemm::default_epilogue(
+ c_buffer, curr_output, curr_m_size, gemm_n_tile_size,
+ output_stride);
+ }
+
+ curr_b_buffer +=
+ b_buffer_block_stride * (gemm_n_tile_size / n_block_size);
+ curr_output += gemm_n_tile_size;
+ }
+ curr_input += gemm_m_tile_size * input_stride;
+ init_output += gemm_m_tile_size * output_stride;
+ }
+ }
+ }
+ }
+}
+
+void cpu_gemm_wna16(
+ const torch::Tensor& input, // [M, K]
+ const torch::Tensor&
+ q_weight, // [N / 16, K * 16 / pack_factor], packed as int32
+ torch::Tensor& output, // [M, N]
+ const torch::Tensor& scales, // [group_num, N]
+ const std::optional&
+ zeros, // [group_num, N / pack_factor], packed as int32
+ const std::optional& g_idx, // [K]
+ const std::optional& bias, // [N]
+ const int64_t pack_factor, const std::string& isa_hint) {
+ using cpu_utils::ISA;
+ TORCH_CHECK_EQ(pack_factor, 8); // only supports 4bits
+ const int32_t a_m_size = input.size(0);
+ const int32_t a_k_size = input.size(1);
+ const int64_t a_m_stride = input.stride(0);
+ const int32_t b_n_size = q_weight.size(0) * 16;
+ TORCH_CHECK_EQ(a_k_size % 32, 0);
+ TORCH_CHECK_EQ(b_n_size % 32, 0);
+ const int32_t group_num = scales.size(0);
+ const int32_t group_size = a_k_size / group_num;
+ TORCH_CHECK_EQ(group_size % 2, 0);
+ const int64_t scales_group_stride = scales.stride(0);
+ const int64_t output_m_stride = output.stride(0);
+
+ bool has_zp = zeros.has_value();
+ bool use_desc_act = g_idx.has_value();
+ TORCH_CHECK(!(has_zp && use_desc_act));
+
+ ISA isa = [&]() {
+ if (isa_hint == "amx") {
+ return ISA::AMX;
+ } else if (isa_hint == "vec") {
+ return ISA::VEC;
+ } else {
+ TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
+ }
+ }();
+
+ int32_t* zeros_ptr = has_zp ? zeros->data_ptr() : nullptr;
+ const int64_t zeros_group_stride = has_zp ? zeros->stride(0) : 0;
+ int32_t* g_idx_ptr = use_desc_act ? g_idx->data_ptr() : nullptr;
+
+ VLLM_DISPATCH_16B_TYPES(input.scalar_type(), "cpu_gemm_wna16", [&]() {
+ if (isa == ISA::AMX) {
+ using gemm_t = cpu_micro_gemm::MicroGemm;
+ if (has_zp) {
+ using dequantizer_t = Dequantizer4b;
+ cpu_gemm_wna16_impl(
+ input.data_ptr(), q_weight.data_ptr(),
+ output.data_ptr(), scales.data_ptr(), zeros_ptr,
+ g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr,
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+ scales_group_stride, zeros_group_stride, group_num, group_size,
+ pack_factor);
+ return;
+ }
+ if (use_desc_act) {
+ using dequantizer_t = Dequantizer4b;
+ cpu_gemm_wna16_impl(
+ input.data_ptr(), q_weight.data_ptr(),
+ output.data_ptr(), scales.data_ptr(), zeros_ptr,
+ g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr,
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+ scales_group_stride, zeros_group_stride, group_num, group_size,
+ pack_factor);
+ return;
+ } else {
+ using dequantizer_t = Dequantizer4b;
+ cpu_gemm_wna16_impl(
+ input.data_ptr(), q_weight.data_ptr(),
+ output.data_ptr(), scales.data_ptr(), zeros_ptr,
+ g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr,
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+ scales_group_stride, zeros_group_stride, group_num, group_size,
+ pack_factor);
+ return;
+ }
+ } else if (isa == ISA::VEC) {
+ using gemm_t = cpu_micro_gemm::MicroGemm;
+ if (has_zp) {
+ using dequantizer_t = Dequantizer4b;
+ cpu_gemm_wna16_impl(
+ input.data_ptr(), q_weight.data_ptr(),
+ output.data_ptr(), scales.data_ptr(), zeros_ptr,
+ g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr,
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+ scales_group_stride, zeros_group_stride, group_num, group_size,
+ pack_factor);
+ return;
+ }
+ if (use_desc_act) {
+ using dequantizer_t = Dequantizer4b;
+ cpu_gemm_wna16_impl(
+ input.data_ptr(), q_weight.data_ptr(),
+ output.data_ptr(), scales.data_ptr(), zeros_ptr,
+ g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr,
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+ scales_group_stride, zeros_group_stride, group_num, group_size,
+ pack_factor);
+ return;
+ } else {
+ using dequantizer_t = Dequantizer4b;
+ cpu_gemm_wna16_impl(
+ input.data_ptr(), q_weight.data_ptr(),
+ output.data_ptr(), scales.data_ptr(), zeros_ptr,
+ g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr,
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+ scales_group_stride, zeros_group_stride, group_num, group_size,
+ pack_factor);
+ return;
+ }
+ }
+ });
+}
diff --git a/csrc/cpu/dnnl_helper.cpp b/csrc/cpu/dnnl_helper.cpp
index 02a8072ccf306..cfb6e78cba9a1 100644
--- a/csrc/cpu/dnnl_helper.cpp
+++ b/csrc/cpu/dnnl_helper.cpp
@@ -396,9 +396,9 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
: DNNLMatMulPrimitiveHandler(
static_cast(args), args.ab_type),
m_size_cache_(nullptr) {
- assert(ab_type_ == dnnl::memory::data_type::f32 ||
- ab_type_ == dnnl::memory::data_type::bf16 ||
- ab_type_ == dnnl::memory::data_type::f16);
+ assert(b_type_ == dnnl::memory::data_type::f32 ||
+ b_type_ == dnnl::memory::data_type::bf16 ||
+ b_type_ == dnnl::memory::data_type::f16);
dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
{b_k_stride_, b_n_stride_});
diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp
new file mode 100644
index 0000000000000..87a019773a895
--- /dev/null
+++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp
@@ -0,0 +1,245 @@
+#ifndef CPU_MICRO_GEMM_AMX_HPP
+#define CPU_MICRO_GEMM_AMX_HPP
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
+
+namespace cpu_micro_gemm {
+namespace {
+// AMX specific
+constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
+constexpr static int64_t AMX_TILE_ROW_NUM = 16;
+constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
+
+typedef struct __tile_config {
+ uint8_t palette_id = 1;
+ uint8_t start_row = 0;
+ uint8_t reserved_0[14] = {0};
+ uint16_t colsb[16] = {0};
+ uint8_t rows[16] = {0};
+} __tilecfg;
+
+// 2-2-4 pattern, for 16 < m <= 32
+// TILE 0, 1: load A matrix, row num should be 16, m - 16
+// TILE 2, 3: load B matrix, row num should be 16
+// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
+// - 16
+template
+class TileGemm224 {
+ public:
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ TORCH_CHECK(false, "Unsupported data type for TileGemm224");
+ }
+
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
+ TORCH_CHECK(false, "Unsupported data type for TileGemm224");
+ }
+};
+
+template <>
+class TileGemm224 {
+ public:
+ using scalar_t = c10::BFloat16;
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
+ c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
+ c10::BFloat16* __restrict__ a_tile_1 = a_ptr + lda * AMX_TILE_ROW_NUM;
+ const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
+
+ // B is always packed as 16 output channels block
+ c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
+ c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
+ const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
+
+ float* __restrict__ c_tile_4 = c_ptr;
+ float* __restrict__ c_tile_5 =
+ c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
+ float* __restrict__ c_tile_6 = c_ptr + AMX_TILE_ROW_NUM * ldc;
+ float* __restrict__ c_tile_7 =
+ c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
+ const int32_t c_tile_stride = ldc * sizeof(float);
+
+ if (accum_c) {
+ _tile_loadd(4, c_tile_4, c_tile_stride);
+ _tile_loadd(5, c_tile_5, c_tile_stride);
+ _tile_loadd(6, c_tile_6, c_tile_stride);
+ _tile_loadd(7, c_tile_7, c_tile_stride);
+ } else {
+ _tile_zero(4);
+ _tile_zero(5);
+ _tile_zero(6);
+ _tile_zero(7);
+ }
+
+ for (int32_t k = 0; k < k_times; ++k) {
+ _tile_loadd(0, a_tile_0, a_tile_stride);
+ _tile_stream_loadd(2, b_tile_2, b_tile_stride);
+ _tile_dpbf16ps(4, 0, 2);
+ _tile_stream_loadd(3, b_tile_3, b_tile_stride);
+ _tile_dpbf16ps(5, 0, 3);
+ _tile_loadd(1, a_tile_1, a_tile_stride);
+ _tile_dpbf16ps(6, 1, 2);
+ _tile_dpbf16ps(7, 1, 3);
+
+ // update ptrs
+ a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
+ a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
+ b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ }
+
+ _tile_stored(4, c_tile_4, c_tile_stride);
+ _tile_stored(5, c_tile_5, c_tile_stride);
+ _tile_stored(6, c_tile_6, c_tile_stride);
+ _tile_stored(7, c_tile_7, c_tile_stride);
+ }
+
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
+ const int32_t m_0 = AMX_TILE_ROW_NUM;
+ const int32_t m_1 = m - AMX_TILE_ROW_NUM;
+ config.rows[0] = m_0;
+ config.rows[1] = m_1;
+ config.rows[2] = AMX_TILE_ROW_NUM;
+ config.rows[3] = AMX_TILE_ROW_NUM;
+ config.rows[4] = m_0;
+ config.rows[5] = m_0;
+ config.rows[6] = m_1;
+ config.rows[7] = m_1;
+ _tile_loadconfig(&config);
+ }
+};
+
+// 1-2-2 pattern, for 0 < m <= 16
+// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
+// m, m
+// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
+// num should be 16
+// TILE 6, 7, (6, 7): store results C matrix, row num should be
+// m
+template
+class TileGemm122 {
+ public:
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ TORCH_CHECK(false, "Unsupported data type for TileGemm122");
+ }
+
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
+ TORCH_CHECK(false, "Unsupported data type for TileGemm122");
+ }
+};
+
+template <>
+class TileGemm122 {
+ public:
+ using scalar_t = c10::BFloat16;
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
+ c10::BFloat16* __restrict__ a_tile_1 =
+ a_ptr + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
+ const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
+
+ c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
+ c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
+ c10::BFloat16* __restrict__ b_tile_4 =
+ b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ c10::BFloat16* __restrict__ b_tile_5 =
+ b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ int64_t b_stride = AMX_TILE_ROW_BYTES;
+
+ float* __restrict__ c_tile_6 = c_ptr;
+ float* __restrict__ c_tile_7 = c_ptr + AMX_TILE_ROW_BYTES / sizeof(float);
+ int64_t c_stride = ldc * sizeof(float);
+
+ const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
+ const int32_t k_group_times = k_times / 2;
+ const bool has_tail = (k_times % 2 == 1);
+
+ if (accum_c) {
+ _tile_loadd(6, c_tile_6, c_stride);
+ _tile_loadd(7, c_tile_7, c_stride);
+ } else {
+ _tile_zero(6);
+ _tile_zero(7);
+ }
+
+ for (int32_t k = 0; k < k_group_times; ++k) {
+ _tile_loadd(0, a_tile_0, a_tile_stride);
+ _tile_stream_loadd(2, b_tile_2, b_stride);
+ _tile_dpbf16ps(6, 0, 2);
+ _tile_stream_loadd(3, b_tile_3, b_stride);
+ _tile_dpbf16ps(7, 0, 3);
+ _tile_loadd(1, a_tile_1, a_tile_stride);
+ _tile_stream_loadd(4, b_tile_4, b_stride);
+ _tile_dpbf16ps(6, 1, 4);
+ _tile_stream_loadd(5, b_tile_5, b_stride);
+ _tile_dpbf16ps(7, 1, 5);
+
+ // update ptrs
+ a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
+ a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
+ b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
+ }
+
+ if (has_tail) {
+ _tile_loadd(0, a_tile_0, a_tile_stride);
+ _tile_stream_loadd(2, b_tile_2, b_stride);
+ _tile_dpbf16ps(6, 0, 2);
+ _tile_stream_loadd(3, b_tile_3, b_stride);
+ _tile_dpbf16ps(7, 0, 3);
+ }
+
+ _tile_stored(6, c_tile_6, c_stride);
+ _tile_stored(7, c_tile_7, c_stride);
+ }
+
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
+ config.rows[0] = m;
+ config.rows[1] = m;
+ config.rows[2] = AMX_TILE_ROW_NUM;
+ config.rows[3] = AMX_TILE_ROW_NUM;
+ config.rows[4] = AMX_TILE_ROW_NUM;
+ config.rows[5] = AMX_TILE_ROW_NUM;
+ config.rows[6] = m;
+ config.rows[7] = m;
+ _tile_loadconfig(&config);
+ }
+};
+} // namespace
+
+// Gemm kernel uses AMX, requires B matrix to be packed
+template
+class MicroGemm {
+ public:
+ static constexpr int32_t MaxMSize = 32;
+ static constexpr int32_t NSize = 32;
+
+ public:
+ MicroGemm() : curr_m_(-1) {
+ vec_op::unroll_loop([&](int i) { amx_tile_config_.colsb[i] = 64; });
+ }
+
+ void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ if (m > AMX_TILE_ROW_NUM) {
+ if (m != curr_m_) {
+ curr_m_ = m;
+ TileGemm224::init_tile_config(m, amx_tile_config_);
+ }
+ TileGemm224::gemm(CPU_MICRO_GEMM_PARAMS);
+ } else {
+ if (m != curr_m_) {
+ curr_m_ = m;
+ TileGemm122::init_tile_config(m, amx_tile_config_);
+ }
+ TileGemm122::gemm(CPU_MICRO_GEMM_PARAMS);
+ }
+ }
+
+ private:
+ alignas(64) __tilecfg amx_tile_config_;
+ int32_t curr_m_;
+};
+
+} // namespace cpu_micro_gemm
+
+#endif
diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp
new file mode 100644
index 0000000000000..784da55a420e5
--- /dev/null
+++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp
@@ -0,0 +1,91 @@
+#ifndef CPU_MICRO_GEMM_IMPL_HPP
+#define CPU_MICRO_GEMM_IMPL_HPP
+#include "cpu/utils.hpp"
+#include "cpu/cpu_types.hpp"
+
+namespace cpu_micro_gemm {
+#define DEFINE_CPU_MICRO_GEMM_PARAMS \
+ scalar_t *__restrict__ a_ptr, scalar_t *__restrict__ b_ptr, \
+ float *__restrict__ c_ptr, const int32_t m, const int32_t k, \
+ const int64_t lda, const int64_t b_n_group_stride, const int64_t ldc, \
+ const bool accum_c
+
+#define CPU_MICRO_GEMM_PARAMS \
+ a_ptr, b_ptr, c_ptr, m, k, lda, b_n_group_stride, ldc, accum_c
+
+template
+class MicroGemm {
+ public:
+ static constexpr int32_t MaxMSize = 16;
+ static constexpr int32_t NSize = 16;
+
+ public:
+ void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ TORCH_CHECK(false, "Unimplemented MicroGemm.");
+ }
+};
+
+template
+FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
+ scalar_t* __restrict__ d_ptr,
+ const int32_t m, const int64_t ldc,
+ const int64_t ldd) {
+ using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t;
+ static_assert(n_size % 16 == 0);
+
+ float* __restrict__ curr_c = c_ptr;
+ scalar_t* __restrict__ curr_d = d_ptr;
+ for (int32_t i = 0; i < m; ++i) {
+ float* __restrict__ curr_c_iter = curr_c;
+ scalar_t* __restrict__ curr_d_iter = curr_d;
+ vec_op::unroll_loop([&](int32_t n_g_idx) {
+ vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
+ scalar_vec_t c_vec(c_vec_fp32);
+ c_vec.save(curr_d_iter);
+ curr_c_iter += 16;
+ curr_d_iter += 16;
+ });
+ curr_c += ldc;
+ curr_d += ldd;
+ }
+}
+
+template
+FORCE_INLINE void bias_epilogue(float* __restrict__ c_ptr,
+ scalar_t* __restrict__ d_ptr,
+ scalar_t* __restrict__ bias_ptr,
+ const int32_t m, const int64_t ldc,
+ const int64_t ldd) {
+ using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t;
+ static_assert(n_size % 16 == 0);
+ constexpr int32_t n_group_num = n_size / 16;
+ static_assert(n_group_num <= 16);
+
+ vec_op::FP32Vec16 bias_vecs[n_group_num];
+ scalar_t* __restrict__ curr_bias = bias_ptr;
+ vec_op::unroll_loop([&](int32_t i) {
+ scalar_vec_t vec(curr_bias);
+ bias_vecs[i] = vec_op::FP32Vec16(vec);
+ curr_bias += 16;
+ });
+
+ float* __restrict__ curr_c = c_ptr;
+ scalar_t* __restrict__ curr_d = d_ptr;
+ for (int32_t i = 0; i < m; ++i) {
+ float* __restrict__ curr_c_iter = curr_c;
+ scalar_t* __restrict__ curr_d_iter = curr_d;
+ vec_op::unroll_loop([&](int32_t n_g_idx) {
+ vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
+ c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
+ scalar_vec_t c_vec(c_vec_fp32);
+ c_vec.save(curr_d_iter);
+ curr_c_iter += 16;
+ curr_d_iter += 16;
+ });
+ curr_c += ldc;
+ curr_d += ldd;
+ }
+}
+} // namespace cpu_micro_gemm
+
+#endif
diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp
new file mode 100644
index 0000000000000..3985c2f2e5fe4
--- /dev/null
+++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp
@@ -0,0 +1,115 @@
+#ifndef CPU_MICRO_GEMM_VEC_HPP
+#define CPU_MICRO_GEMM_VEC_HPP
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
+
+namespace cpu_micro_gemm {
+namespace {
+// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
+template
+class TileGemm82 {
+ public:
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ switch (m) {
+ case 1:
+ gemm_micro<1>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 2:
+ gemm_micro<2>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 3:
+ gemm_micro<3>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 4:
+ gemm_micro<4>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 5:
+ gemm_micro<5>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 6:
+ gemm_micro<6>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 7:
+ gemm_micro<7>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ case 8:
+ gemm_micro<8>(CPU_MICRO_GEMM_PARAMS);
+ break;
+ }
+ }
+
+ template
+ static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ static_assert(0 < M <= 8);
+ using load_vec_t = typename cpu_utils::VecTypeTrait::vec_t;
+
+ scalar_t* __restrict__ curr_b_0 = b_ptr;
+ scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
+ float* __restrict__ curr_c_0 = c_ptr;
+ float* __restrict__ curr_c_1 = c_ptr + 16;
+
+ vec_op::FP32Vec16 c_regs[M * 2];
+ if (accum_c) {
+ float* __restrict__ curr_m_c_0 = curr_c_0;
+ float* __restrict__ curr_m_c_1 = curr_c_1;
+ vec_op::unroll_loop([&](int32_t i) {
+ c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
+ c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
+
+ // update
+ curr_m_c_0 += ldc;
+ curr_m_c_1 += ldc;
+ });
+ }
+
+ scalar_t* __restrict__ curr_a = a_ptr;
+ for (int32_t k_idx = 0; k_idx < k; ++k_idx) {
+ load_vec_t b_0_reg(curr_b_0);
+ vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
+ load_vec_t b_1_reg(curr_b_1);
+ vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
+
+ scalar_t* __restrict__ curr_m_a = curr_a;
+ vec_op::unroll_loop([&](int32_t i) {
+ scalar_t v = *curr_m_a;
+ load_vec_t a_reg_original(v);
+ vec_op::FP32Vec16 a_reg(a_reg_original);
+ c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
+ c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
+
+ // update
+ curr_m_a += lda;
+ });
+
+ // update
+ curr_a += 1;
+ curr_b_0 += 16;
+ curr_b_1 += 16;
+ }
+
+ vec_op::unroll_loop([&](int32_t i) {
+ c_regs[i * 2].save(curr_c_0);
+ c_regs[i * 2 + 1].save(curr_c_1);
+
+ // update
+ curr_c_0 += ldc;
+ curr_c_1 += ldc;
+ });
+ }
+};
+} // namespace
+
+// Gemm kernel uses vector instructions, requires B matrix to be packed
+template
+class MicroGemm {
+ public:
+ static constexpr int32_t MaxMSize = 8;
+ static constexpr int32_t NSize = 32;
+
+ public:
+ void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+ TileGemm82::gemm(CPU_MICRO_GEMM_PARAMS);
+ }
+};
+} // namespace cpu_micro_gemm
+
+#endif
diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp
index 9fefd88cd9b08..e0e3ef71b485f 100644
--- a/csrc/cpu/torch_bindings.cpp
+++ b/csrc/cpu/torch_bindings.cpp
@@ -103,6 +103,13 @@ void cpu_attention_with_kv_cache(
// Note: just for avoiding importing errors
void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
+void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight,
+ torch::Tensor& output, const torch::Tensor& scales,
+ const std::optional& zeros,
+ const std::optional& g_idx,
+ const std::optional& bias,
+ const int64_t pack_factor, const std::string& isa_hint);
+
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
@@ -165,7 +172,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Quantization
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
defined(__powerpc64__)
- at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
// Helper function to release oneDNN handlers
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
&release_dnnl_matmul_handler);
@@ -201,15 +207,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Compute int8 quantized tensor for given scaling factor.
ops.def(
"static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
- "Tensor? azp) -> ()",
- {stride_tag});
+ "Tensor? azp) -> ()");
ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
// Compute int8 quantized tensor and scaling factor
ops.def(
"dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
- "Tensor!? azp) -> ()",
- {stride_tag});
+ "Tensor!? azp) -> ()");
ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
&dynamic_scaled_int8_quant);
#endif
@@ -283,6 +287,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("static_scaled_fp8_quant() -> ()", placeholder_op);
ops.def("dynamic_scaled_fp8_quant() -> ()", placeholder_op);
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
+
+ // WNA16
+#if defined(__AVX512F__)
+ ops.def(
+ "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
+ "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
+ "pack_factor, str isa_hint) -> ()");
+ ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16);
+#endif
}
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
diff --git a/csrc/cpu/utils.cpp b/csrc/cpu/utils.cpp
index c5a48352e3089..5199ba2af024f 100644
--- a/csrc/cpu/utils.cpp
+++ b/csrc/cpu/utils.cpp
@@ -45,31 +45,54 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
// Memory node binding
if (numa_available() != -1) {
int mem_node_id = numa_node_of_cpu(omp_cpu_ids.front());
- // Verify all CPUs are on the same NUMA node
- for (size_t i = 1; i < omp_cpu_ids.size(); ++i) {
- int node_id = numa_node_of_cpu(omp_cpu_ids[i]);
- TORCH_CHECK(node_id == mem_node_id, "CPU ", omp_cpu_ids[i],
- " is on NUMA node ", node_id, ", but CPU ",
- omp_cpu_ids.front(), " is on NUMA node ", mem_node_id,
- ". All CPUs should be on the same NUMA node for optimal "
- "performance. Memory will be bound to NUMA node ",
- mem_node_id, ".");
+ std::set node_ids;
+ for (const auto& cpu_id : omp_cpu_ids) {
+ int node_id = numa_node_of_cpu(cpu_id);
+ if (node_id != -1) {
+ node_ids.insert(node_id);
+ }
+ TORCH_WARN(node_id == mem_node_id, "CPU ", cpu_id, " is on NUMA node ",
+ node_id, ", but CPU ", omp_cpu_ids.front(),
+ " is on NUMA node ", mem_node_id,
+ ". All CPUs should be on the same NUMA node for optimal "
+ "performance. Memory will be bound to NUMA node ",
+ mem_node_id, ".");
}
- bitmask* mask = numa_parse_nodestring(std::to_string(mem_node_id).c_str());
- bitmask* src_mask = numa_get_membind();
+ // Concatenate all node_ids into a single comma-separated string
+ if (!node_ids.empty()) {
+ std::string node_ids_str;
+ for (const int node_id : node_ids) {
+ if (!node_ids_str.empty()) {
+ node_ids_str += ",";
+ }
+ node_ids_str += std::to_string(node_id);
+ }
- int pid = getpid();
+ bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
+ bitmask* src_mask = numa_get_membind();
- // move all existing pages to the specified numa node.
- *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
- int page_num = numa_migrate_pages(pid, src_mask, mask);
- if (page_num == -1) {
- TORCH_WARN("numa_migrate_pages failed. errno: " + std::to_string(errno));
+ int pid = getpid();
+
+ if (mask && src_mask) {
+ // move all existing pages to the specified numa node.
+ *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
+ int page_num = numa_migrate_pages(pid, src_mask, mask);
+ if (page_num == -1) {
+ TORCH_WARN("numa_migrate_pages failed. errno: " +
+ std::to_string(errno));
+ }
+
+ // restrict memory allocation node.
+ numa_set_membind(mask);
+ numa_set_strict(1);
+
+ numa_free_nodemask(mask);
+ numa_free_nodemask(src_mask);
+ } else {
+ TORCH_WARN("numa_parse_nodestring or numa_get_membind failed. errno: " +
+ std::to_string(errno));
+ }
}
-
- // restrict memory allocation node.
- numa_set_membind(mask);
- numa_set_strict(1);
}
// OMP threads binding
diff --git a/csrc/cpu/utils.hpp b/csrc/cpu/utils.hpp
new file mode 100644
index 0000000000000..d3def306b8069
--- /dev/null
+++ b/csrc/cpu/utils.hpp
@@ -0,0 +1,73 @@
+#ifndef UTILS_HPP
+#define UTILS_HPP
+
+#include
+#include
+#include
+#include
+
+#if defined(__APPLE__)
+ #include
+#endif
+
+#include "cpu_types.hpp"
+
+namespace cpu_utils {
+enum class ISA { AMX, VEC };
+
+template
+struct VecTypeTrait {
+ using vec_t = void;
+};
+
+template <>
+struct VecTypeTrait {
+ using vec_t = vec_op::FP32Vec16;
+};
+
+#if !defined(__aarch64__) || defined(ARM_BF16_SUPPORT)
+template <>
+struct VecTypeTrait {
+ using vec_t = vec_op::BF16Vec16;
+};
+#endif
+
+template <>
+struct VecTypeTrait {
+ using vec_t = vec_op::FP16Vec16;
+};
+
+struct Counter {
+ std::atomic counter;
+ char _padding[56];
+
+ Counter() : counter(0) {}
+
+ void reset_counter() { counter.store(0); }
+
+ int64_t acquire_counter() { return counter++; }
+};
+
+inline int64_t get_l2_size() {
+ static int64_t size = []() {
+#if defined(__APPLE__)
+ // macOS doesn't have _SC_LEVEL2_CACHE_SIZE. Use sysctlbyname.
+ int64_t l2_cache_size = 0;
+ size_t len = sizeof(l2_cache_size);
+ if (sysctlbyname("hw.l2cachesize", &l2_cache_size, &len, NULL, 0) == 0 &&
+ l2_cache_size > 0) {
+ return l2_cache_size >> 1; // use 50% of L2 cache
+ }
+ // Fallback if sysctlbyname fails
+ return 128LL * 1024 >> 1; // use 50% of 128KB
+#else
+ long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
+ assert(l2_cache_size != -1);
+ return l2_cache_size >> 1; // use 50% of L2 cache
+#endif
+ }();
+ return size;
+}
+} // namespace cpu_utils
+
+#endif
diff --git a/csrc/cuda_view.cu b/csrc/cuda_view.cu
index 938bd4ab7fc62..9853fc942bab7 100644
--- a/csrc/cuda_view.cu
+++ b/csrc/cuda_view.cu
@@ -22,15 +22,10 @@ torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
auto strides = cpu_tensor.strides();
auto options = cpu_tensor.options().device(torch::kCUDA);
- // from_blob signature: from_blob(void *data, IntArrayRef sizes, ..., Deleter,
- // const TensorOptions &) Provide a no-op deleter. The CPU tensor holds the
- // memory, so we don't free it here.
- auto deleter = [](void*) {
- // no-op, since the memory is owned by the original CPU tensor
- };
-
+ // use default no-op deleter, since the memory is owned by the original CPU
+ // tensor
torch::Tensor cuda_tensor =
- torch::from_blob(device_ptr, sizes, strides, deleter, options);
+ torch::from_blob(device_ptr, sizes, strides, options);
TORCH_CHECK(cuda_tensor.device().is_cuda(),
"Resulting tensor is not on CUDA device");
diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h
index 9ae0ed975edde..e1d131e4a7851 100644
--- a/csrc/dispatch_utils.h
+++ b/csrc/dispatch_utils.h
@@ -117,3 +117,24 @@
break; \
} \
}
+
+#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \
+ switch (NUM_DIMS) { \
+ case 2: { \
+ constexpr int tensor_rank = 2; \
+ __VA_ARGS__(); \
+ break; \
+ } \
+ case 3: { \
+ constexpr int tensor_rank = 3; \
+ __VA_ARGS__(); \
+ break; \
+ } \
+ case 4: { \
+ constexpr int tensor_rank = 4; \
+ __VA_ARGS__(); \
+ break; \
+ } \
+ default: \
+ TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \
+ }
diff --git a/csrc/layernorm_kernels.cu b/csrc/layernorm_kernels.cu
index 48771e4b3aff9..dfc67b933ccae 100644
--- a/csrc/layernorm_kernels.cu
+++ b/csrc/layernorm_kernels.cu
@@ -10,16 +10,38 @@
namespace vllm {
// TODO(woosuk): Further optimize this kernel.
-template
+template
__global__ void rms_norm_kernel(
- scalar_t* __restrict__ out, // [..., hidden_size]
- const scalar_t* __restrict__ input, // [..., hidden_size]
- const int64_t input_stride,
+ scalar_t* __restrict__ out, // [..., hidden_size]
+ const scalar_t* __restrict__ input, // [..., hidden_size]
+ const int64_t input_stride_d2, // input.stride(-2)
+ const int64_t input_stride_d3, // input.stride(-3)
+ 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]
const float epsilon, const int num_tokens, const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
- const scalar_t* input_row = input + blockIdx.x * input_stride;
+ const scalar_t* input_row;
+ if constexpr (NUM_DIMS == 2) {
+ // 2D for layernorm normal case [batch_size, hidden]
+ input_row = input + blockIdx.x * input_stride_d2;
+ } 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;
+ } 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);
+ int remaining = blockIdx.x % (input_shape_d3 * input_shape_d2);
+ int seq_idx = remaining / input_shape_d2;
+ 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;
+ }
auto vec_op = [&variance](const vec_n_t& vec) {
#pragma unroll
@@ -164,38 +186,44 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]
torch::Tensor& weight, // [hidden_size]
double epsilon) {
TORCH_CHECK(out.is_contiguous());
+ if (input.stride(-1) != 1) {
+ input = input.contiguous();
+ }
TORCH_CHECK(input.stride(-1) == 1);
TORCH_CHECK(weight.is_contiguous());
int hidden_size = input.size(-1);
- // We cannot just use `input.stride(-2)` if the tensor is not row-major.
- // Instead, we use a 2d view to get the second-innermost stride.
- // That way the dimensions (except the last one) can be arbitrarily permuted.
- torch::Tensor input_view = input.view({-1, hidden_size});
-
- int num_tokens = input_view.numel() / hidden_size;
- int64_t input_stride = input_view.stride(-2);
+ int num_tokens = input.numel() / hidden_size;
+ int num_dims = input.dim();
+ int64_t input_stride_d2 = input.stride(-2);
+ int64_t input_stride_d3 = (num_dims >= 3) ? input.stride(-3) : 0;
+ int64_t input_stride_d4 = (num_dims >= 4) ? input.stride(-4) : 0;
+ int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0;
+ int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
// For large num_tokens, use smaller blocks to increase SM concurrency.
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 grid(num_tokens);
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input_view));
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
- VLLM_DISPATCH_FLOATING_TYPES(
- input_view.scalar_type(), "rms_norm_kernel", [&] {
- const int calculated_vec_size =
- std::gcd(16 / sizeof(scalar_t), hidden_size);
- const int block_size =
- std::min(hidden_size / calculated_vec_size, max_block_size);
- dim3 block(block_size);
- VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
- vllm::rms_norm_kernel<<>>(
- out.data_ptr(), input_view.data_ptr(),
- input_stride, weight.data_ptr(), epsilon, num_tokens,
- hidden_size);
- });
+ VLLM_DISPATCH_RANK234(num_dims, [&] {
+ VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
+ const int calculated_vec_size =
+ std::gcd(16 / sizeof(scalar_t), hidden_size);
+ const int block_size =
+ std::min(hidden_size / calculated_vec_size, max_block_size);
+ dim3 block(block_size);
+ VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+ vllm::rms_norm_kernel
+ <<>>(
+ out.data_ptr(), input.data_ptr(),
+ input_stride_d2, input_stride_d3, input_stride_d4,
+ input_shape_d2, input_shape_d3, weight.data_ptr(),
+ epsilon, num_tokens, hidden_size);
});
+ });
+ });
}
#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp
index c3ae06a30e3e8..5af74c2c2a6b0 100644
--- a/csrc/torch_bindings.cpp
+++ b/csrc/torch_bindings.cpp
@@ -20,18 +20,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
//
- // The default behavior in PyTorch 2.6 was changed to "requires_contiguous",
- // so we need
- // to override this for many GEMMs with the following tag. Otherwise,
- // torch.compile will force all input tensors to be contiguous(), which
- // will break many custom ops that require column-major weight matrices.
- // This was a bug and PyTorch 2.7 has since fixed this.
-#if TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 6
- #define stride_tag at::Tag::needs_fixed_stride_order
-#else
- #define stride_tag
-#endif
-
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s,"
@@ -241,15 +229,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Quantized GEMM for AWQ.
ops.def(
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
- "Tensor _zeros, SymInt split_k_iters) -> Tensor",
- {stride_tag});
+ "Tensor _zeros, SymInt split_k_iters) -> Tensor");
ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
// Dequantization for AWQ.
ops.def(
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
- "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor",
- {stride_tag});
+ "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
// Note about marlin kernel 'workspace' arguments:
@@ -271,8 +257,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
"Tensor b_scales, Tensor workspace, "
"int b_q_type, "
- "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor",
- {stride_tag});
+ "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor");
// conditionally compiled so impl in source file
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
@@ -298,8 +283,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" Tensor? channel_scales,"
" Tensor? token_scales,"
" str? schedule"
- ") -> Tensor",
- {stride_tag});
+ ") -> Tensor");
ops.def(
"machete_prepack_B("
" Tensor B,"
@@ -319,8 +303,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
"g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
- "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
- {stride_tag});
+ "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
// conditionally compiled so impl registration is in source file
// gptq_marlin repack from GPTQ.
@@ -346,8 +329,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" Tensor token_scales,"
" ScalarType? out_type,"
" str? maybe_schedule"
- ") -> Tensor",
- {stride_tag});
+ ") -> Tensor");
// pack scales
ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
// encode and reorder weight matrix
@@ -394,24 +376,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
" Tensor block_scale_a, Tensor block_scale_b,"
- " Tensor alpha) -> ()",
- {stride_tag});
+ " Tensor alpha) -> ()");
ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
// cutlass blockwise scaledgroup GEMM
ops.def(
"cutlass_blockwise_scaled_grouped_mm(Tensor! output, Tensor a, Tensor b, "
"Tensor scales_a, Tensor scales_b, "
- "Tensor problem_sizes, Tensor expert_offsets) -> ()",
- {stride_tag});
+ "Tensor problem_sizes, Tensor expert_offsets) -> ()");
// conditionally compiled so impl registration is in source file
// cutlass nvfp4 block scaled group GEMM
ops.def(
"cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
- " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()",
- {stride_tag});
+ " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
// conditionally compiled so impl registration is in source file
// CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
@@ -419,8 +398,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"cutlass_scaled_mm(Tensor! out, Tensor a,"
" Tensor b, Tensor a_scales,"
- " Tensor b_scales, Tensor? bias) -> ()",
- {stride_tag});
+ " Tensor b_scales, Tensor? bias) -> ()");
ops.impl("cutlass_scaled_mm", torch::kCUDA, &cutlass_scaled_mm);
// CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
@@ -429,8 +407,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
" Tensor b, Tensor a_scales,"
" Tensor b_scales, Tensor azp_adj,"
- " Tensor? azp, Tensor? bias) -> ()",
- {stride_tag});
+ " Tensor? azp, Tensor? bias) -> ()");
ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);
// Check if cutlass scaled_mm is supported for CUDA devices of the given
@@ -449,8 +426,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
" Tensor problem_sizes, Tensor a_strides, "
" Tensor b_strides, Tensor c_strides, bool per_act_token, "
- " bool per_out_ch) -> ()",
- {stride_tag});
+ " bool per_out_ch) -> ()");
ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);
// A function that computes data required to run fused MoE with w8a8 grouped
@@ -464,8 +440,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" Tensor! problem_sizes1, Tensor! problem_sizes2, "
" Tensor! input_permutation, "
" Tensor! output_permutation, int num_experts, "
- " int n, int k, Tensor? blockscale_offsets) -> ()",
- {stride_tag});
+ " int n, int k, Tensor? blockscale_offsets) -> "
+ "()");
ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);
// A function that computes problem sizes for each expert's multiplication
@@ -476,8 +452,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" Tensor! problem_sizes1, "
" Tensor! problem_sizes2, "
" int num_experts, int n, int k, "
- " Tensor? blockscale_offsets) -> ()",
- {stride_tag});
+ " Tensor? blockscale_offsets) -> ()");
ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,
&get_cutlass_moe_mm_problem_sizes);
@@ -492,8 +467,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" Tensor! problem_sizes2, "
" Tensor expert_num_tokens, "
" int num_local_experts, int padded_m, "
- " int n, int k) -> ()",
- {stride_tag});
+ " int n, int k) -> ()");
ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA,
&get_cutlass_pplx_moe_mm_data);
@@ -517,8 +491,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
" Tensor bt_nzs,"
" Tensor bt_meta, Tensor a_scales,"
- " Tensor b_scales, Tensor? bias) -> ()",
- {stride_tag});
+ " Tensor b_scales, Tensor? bias) -> ()");
ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);
// CUTLASS sparse matrix compressor
@@ -567,8 +540,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
"Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
"use_v2_format, int bit) "
- "-> Tensor",
- {stride_tag});
+ "-> Tensor");
ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
// Post processing for GPTQ.
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 964700e2a43ac..1b937bbc1225e 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -56,7 +56,6 @@ ARG UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
# PyTorch provides its own indexes for standard and nightly builds
ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl
-ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL=https://download.pytorch.org/whl/nightly
# PIP supports multiple authentication schemes, including keyring
# By parameterizing the PIP_KEYRING_PROVIDER variable and setting it to
@@ -86,7 +85,7 @@ ARG GET_PIP_URL
RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
&& echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
&& apt-get update -y \
- && apt-get install -y ccache software-properties-common git curl sudo python3-pip \
+ && apt-get install -y ccache software-properties-common git curl sudo python3-pip libibverbs-dev \
&& curl -LsSf https://astral.sh/uv/install.sh | sh \
&& $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
&& rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
@@ -98,7 +97,6 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL
-ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
# Activate virtual environment and add uv to PATH
@@ -226,6 +224,22 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
fi
+# Install DeepGEMM from source
+ARG DEEPGEMM_GIT_REF
+COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
+RUN --mount=type=cache,target=/root/.cache/uv \
+ VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} --wheel-dir /tmp/deepgemm/dist
+
+# Ensure the wheel dir exists so later-stage COPY won't fail when DeepGEMM is skipped
+RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
+
+COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
+# Install EP kernels(pplx-kernels and DeepEP)
+RUN --mount=type=cache,target=/root/.cache/uv \
+ export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
+ /tmp/install_python_libraries.sh /tmp/ep_kernels_workspace wheel && \
+ find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
+
# Check the size of the wheel if RUN_WHEEL_CHECK is true
COPY .buildkite/check-wheel-size.py check-wheel-size.py
# sync the default value with .buildkite/check-wheel-size.py
@@ -291,7 +305,7 @@ RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
&& echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
&& apt-get update -y \
- && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \
+ && apt-get install -y software-properties-common curl sudo python3-pip \
&& apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
&& if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
@@ -317,7 +331,6 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL
-ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
# Install uv for faster pip installs
@@ -337,20 +350,6 @@ ENV UV_LINK_MODE=copy
# or future versions of triton.
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
-# arm64 (GH200) build follows the practice of "use existing pytorch" build,
-# we need to install torch and torchvision from the nightly builds first,
-# pytorch will not appear as a vLLM dependency in all of the following steps
-# after this step
-RUN --mount=type=cache,target=/root/.cache/uv \
- if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
- uv pip install --system \
- --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
- "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319" ; \
- uv pip install --system \
- --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
- --pre pytorch_triton==3.3.0+gitab727c40 ; \
- fi
-
# Install vllm wheel first, so that torch etc will be installed.
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
--mount=type=cache,target=/root/.cache/uv \
@@ -373,36 +372,32 @@ RUN --mount=type=cache,target=/root/.cache/uv \
. /etc/environment && \
uv pip list
-# Even when we build Flashinfer with AOT mode, there's still
-# some issues w.r.t. JIT compilation. Therefore we need to
-# install build dependencies for JIT compilation.
-# TODO: Remove this once FlashInfer AOT wheel is fixed
-COPY requirements/build.txt requirements/build.txt
+# Install deepgemm wheel that has been built in the `build` stage
RUN --mount=type=cache,target=/root/.cache/uv \
- uv pip install --system -r requirements/build.txt \
+ --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
+ sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \
+ uv pip install --system /tmp/deepgemm/dist/*.whl; \
+ else \
+ echo "No DeepGEMM wheels to install; skipping."; \
+ fi'
+
+# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH (https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/.ci/manywheel/build_cuda.sh#L141C14-L141C36)
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+
+# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage
+RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
+ --mount=type=cache,target=/root/.cache/uv \
+ uv pip install --system ep_kernels/dist/*.whl --verbose \
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
-# Install DeepGEMM from source
-ARG DEEPGEMM_GIT_REF
-COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
-RUN --mount=type=cache,target=/root/.cache/uv \
- VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"}
-
-COPY tools/install_gdrcopy.sh install_gdrcopy.sh
-RUN set -eux; \
+RUN --mount=type=bind,source=tools/install_gdrcopy.sh,target=/tmp/install_gdrcopy.sh,ro \
+ set -eux; \
case "${TARGETPLATFORM}" in \
linux/arm64) UUARCH="aarch64" ;; \
linux/amd64) UUARCH="x64" ;; \
*) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
esac; \
- ./install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \
- rm ./install_gdrcopy.sh
-
-# Install EP kernels(pplx-kernels and DeepEP)
-COPY tools/ep_kernels/install_python_libraries.sh install_python_libraries.sh
-ENV CUDA_HOME=/usr/local/cuda
-RUN export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0a 10.0a+PTX}" \
- && bash install_python_libraries.sh
+ /tmp/install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"
# CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
# return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
@@ -432,6 +427,11 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
# Use copy mode to avoid hardlink failures with Docker cache mounts
ENV UV_LINK_MODE=copy
+RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
+ && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
+ && apt-get update -y \
+ && apt-get install -y git
+
# install development dependencies (for testing)
RUN --mount=type=cache,target=/root/.cache/uv \
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
@@ -472,12 +472,11 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500
-COPY requirements/kv_connectors.txt requirements/kv_connectors.txt
-
# install additional dependencies for openai api server
RUN --mount=type=cache,target=/root/.cache/uv \
+ --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
- uv pip install --system -r requirements/kv_connectors.txt; \
+ uv pip install --system -r /tmp/kv_connectors.txt; \
fi; \
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
BITSANDBYTES_VERSION="0.42.0"; \
diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu
index 4c961defaeda2..eb3807ef0ca4e 100644
--- a/docker/Dockerfile.cpu
+++ b/docker/Dockerfile.cpu
@@ -37,6 +37,7 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
&& curl -LsSf https://astral.sh/uv/install.sh | sh
+ENV CC=/usr/bin/gcc-12 CXX=/usr/bin/g++-12
ENV CCACHE_DIR=/root/.cache/ccache
ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache
@@ -122,6 +123,15 @@ WORKDIR /workspace/vllm
RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
cp requirements/test.in requirements/cpu-test.in && \
sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
+ remove_packages_not_supported_on_aarch64() { \
+ case "$(uname -m)" in \
+ aarch64|arm64) \
+ sed -i '/decord/d' requirements/cpu-test.in; \
+ sed -i '/terratorch/d' requirements/cpu-test.in; \
+ ;; \
+ esac; \
+ }; \
+ remove_packages_not_supported_on_aarch64 && \
sed -i 's/^torch==.*/torch==2.8.0/g' requirements/cpu-test.in && \
sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch
index b88b9c4992200..d663c82c3885e 100644
--- a/docker/Dockerfile.nightly_torch
+++ b/docker/Dockerfile.nightly_torch
@@ -76,34 +76,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/common.txt
-# must put before installing xformers, so it can install the correct version of xfomrers.
-ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
-ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
-
-# Build xformers with cuda and torch nightly
-# following official xformers guidance: https://github.com/facebookresearch/xformers#build
-# todo(elainewy): cache xformers build result for faster build
-ARG max_jobs=16
-ENV MAX_JOBS=${max_jobs}
-ARG XFORMERS_COMMIT=f2de641ef670510cadab099ce6954031f52f191c
-
-ENV CCACHE_DIR=/root/.cache/ccache
-RUN --mount=type=cache,target=/root/.cache/ccache \
- --mount=type=cache,target=/root/.cache/uv \
- echo 'git clone xformers...' \
- && git clone https://github.com/facebookresearch/xformers.git --recursive \
- && cd xformers \
- && git checkout ${XFORMERS_COMMIT} \
- && git submodule update --init --recursive \
- && echo 'finish git clone xformers...' \
- && rm -rf build \
- && python3 setup.py bdist_wheel --dist-dir=../xformers-dist --verbose \
- && cd .. \
- && rm -rf xformers
-
-RUN --mount=type=cache,target=/root/.cache/uv \
- uv pip install --system xformers-dist/*.whl --verbose
-
# build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
# track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
@@ -233,11 +205,6 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/vllm
--mount=type=cache,target=/root/.cache/uv \
uv pip install --system vllm-dist/*.whl --verbose
-# install xformers again for the new environment
-RUN --mount=type=bind,from=base,src=/workspace/xformers-dist,target=/vllm-workspace/xformers-dist \
- --mount=type=cache,target=/root/.cache/uv \
- uv pip install --system /vllm-workspace/xformers-dist/*.whl --verbose
-
ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
# install package for build flashinfer
@@ -307,7 +274,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/nightly_torch_test.txt
# Logging to confirm the torch versions
-RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'
+RUN pip freeze | grep -E 'torch|vllm|flashinfer'
# Logging to confirm all the packages are installed
RUN pip freeze
diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le
index ad9eae94b83dd..b16bea3607d2f 100644
--- a/docker/Dockerfile.ppc64le
+++ b/docker/Dockerfile.ppc64le
@@ -8,8 +8,8 @@ FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS openbl
ARG MAX_JOBS
ARG OPENBLAS_VERSION=0.3.30
-RUN microdnf install -y dnf && dnf install -y gcc-toolset-13 make wget unzip \
- && source /opt/rh/gcc-toolset-13/enable \
+RUN microdnf install -y dnf && dnf install -y gcc-toolset-14 make wget unzip \
+ && source /opt/rh/gcc-toolset-14/enable \
&& wget https://github.com/OpenMathLib/OpenBLAS/releases/download/v$OPENBLAS_VERSION/OpenBLAS-$OPENBLAS_VERSION.zip \
&& unzip OpenBLAS-$OPENBLAS_VERSION.zip \
&& cd OpenBLAS-$OPENBLAS_VERSION \
@@ -57,7 +57,7 @@ COPY --from=openblas-builder /tmp/control /dev/null
RUN --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \
dnf install -y openssl-devel \
&& dnf install -y \
- git tar gcc-toolset-13 automake libtool \
+ git tar gcc-toolset-14 automake libtool \
pkgconfig xsimd zeromq-devel kmod findutils protobuf* \
libtiff-devel libjpeg-devel zlib-devel freetype-devel libwebp-devel \
harfbuzz-devel libraqm-devel libimagequant-devel libxcb-devel \
@@ -84,7 +84,7 @@ ARG _GLIBCXX_USE_CXX11_ABI=1
ARG OPENBLAS_VERSION=0.3.30
RUN --mount=type=cache,target=/root/.cache/uv \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
git clone --recursive https://github.com/pytorch/pytorch.git -b v${TORCH_VERSION} && \
cd pytorch && \
uv pip install -r requirements.txt && \
@@ -97,7 +97,7 @@ ARG TORCHVISION_VERSION=0.22.0
ARG TORCHVISION_USE_NVJPEG=0
ARG TORCHVISION_USE_FFMPEG=0
RUN --mount=type=cache,target=/root/.cache/uv \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
git clone --recursive https://github.com/pytorch/vision.git -b v${TORCHVISION_VERSION} && \
cd vision && \
MAX_JOBS=${MAX_JOBS:-$(nproc)} \
@@ -113,7 +113,7 @@ ARG USE_ROCM=0
ARG USE_CUDA=0
ARG TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1
RUN --mount=type=cache,target=/root/.cache/uv \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} && \
cd audio && \
MAX_JOBS=${MAX_JOBS:-$(nproc)} \
@@ -130,7 +130,7 @@ ARG MAX_JOBS
ARG PYARROW_PARALLEL
ARG PYARROW_VERSION=21.0.0
RUN --mount=type=cache,target=/root/.cache/uv \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
git clone --recursive https://github.com/apache/arrow.git -b apache-arrow-${PYARROW_VERSION} && \
cd arrow/cpp && \
mkdir build && cd build && \
@@ -162,7 +162,7 @@ ARG OPENCV_VERSION=86
ARG OPENCV_PATCH=97f3f39
ARG ENABLE_HEADLESS=1
RUN --mount=type=cache,target=/root/.cache/uv \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
git clone --recursive https://github.com/opencv/opencv-python.git -b ${OPENCV_VERSION} && \
cd opencv-python && \
sed -i -E -e 's/"setuptools.+",/"setuptools",/g' pyproject.toml && \
@@ -196,7 +196,7 @@ ARG MAX_JOBS
ARG NUMBA_VERSION=0.61.2
# Clone all required dependencies
-RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset-13/enable && export PATH=$PATH:/usr/lib64/llvm15/bin && \
+RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset-14/enable && export PATH=$PATH:/usr/lib64/llvm15/bin && \
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
cd ./numba && \
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
@@ -211,6 +211,9 @@ RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset
FROM base-builder AS vllmcache-builder
+ENV LLVM_CONFIG=/usr/lib64/llvm15/bin/llvm-config
+ENV PATH=/usr/lib64/llvm15/bin:$PATH
+
COPY --from=torch-builder /tmp/control /dev/null
COPY --from=arrow-builder /tmp/control /dev/null
COPY --from=cv-builder /tmp/control /dev/null
@@ -225,10 +228,13 @@ ARG GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1
RUN --mount=type=cache,target=/root/.cache/uv \
dnf install llvm15 llvm15-devel -y && \
rpm -ivh --nodeps https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-16.el9.ppc64le.rpm && \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
git clone https://github.com/huggingface/xet-core.git && cd xet-core/hf_xet/ && \
uv pip install maturin && \
uv build --wheel --out-dir /hf_wheels/
+
+ENV CXXFLAGS="-fno-lto -Wno-error=free-nonheap-object" \
+ CFLAGS="-fno-lto"
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=torch-builder,source=/torchwheels/,target=/torchwheels/,ro \
--mount=type=bind,from=arrow-builder,source=/arrowwheels/,target=/arrowwheels/,ro \
@@ -236,7 +242,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=numa-builder,source=/numactl/,target=/numactl/,rw \
--mount=type=bind,from=numba-builder,source=/numbawheels/,target=/numbawheels/,ro \
--mount=type=bind,src=.,dst=/src/,rw \
- source /opt/rh/gcc-toolset-13/enable && \
+ source /opt/rh/gcc-toolset-14/enable && \
export PATH=$PATH:/usr/lib64/llvm15/bin && \
uv pip install /opencvwheels/*.whl /arrowwheels/*.whl /torchwheels/*.whl /numbawheels/*.whl && \
sed -i -e 's/.*torch.*//g' /src/pyproject.toml /src/requirements/*.txt && \
@@ -260,7 +266,7 @@ FROM base-builder AS lapack-builder
ARG MAX_JOBS
ARG LAPACK_VERSION=3.12.1
RUN git clone --recursive https://github.com/Reference-LAPACK/lapack.git -b v${LAPACK_VERSION} \
- && cd lapack && source /opt/rh/gcc-toolset-13/enable \
+ && cd lapack && source /opt/rh/gcc-toolset-14/enable \
&& cmake -B build -S . \
&& cmake --build build -j ${MAX_JOBS:-$(nproc)}
@@ -299,7 +305,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \
rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \
microdnf install --nodocs -y \
- libomp tar findutils openssl llvm15 llvm15-devel \
+ libomp libicu tar findutils openssl llvm15 llvm15-devel \
pkgconfig xsimd g++ gcc-fortran libsndfile \
libtiff libjpeg openjpeg2 zlib zeromq \
freetype lcms2 libwebp tcl tk utf8proc \
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index 731a97d93da1f..42466d1801cf6 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -7,6 +7,8 @@ FROM ${BASE_IMAGE} AS base
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
+ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1
# Install some basic utilities
RUN apt-get update -q -y && apt-get install -q -y \
@@ -121,8 +123,6 @@ COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
-ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
-ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1
ENV TOKENIZERS_PARALLELISM=false
# ENV that can improve safe tensor loading, and end-to-end time
diff --git a/docs/.nav.yml b/docs/.nav.yml
index 3151ea0e2ec22..c8bf00efb2370 100644
--- a/docs/.nav.yml
+++ b/docs/.nav.yml
@@ -24,14 +24,16 @@ nav:
- deployment/integrations
- Training: training
- Configuration:
- - configuration/README.md
- configuration/*
+ - TPU: https://docs.vllm.ai/projects/tpu/en/latest/
- Models:
- models/supported_models.md
- models/generative_models.md
- models/pooling_models.md
- models/extensions
- - Hardware Supported Models: models/hardware_supported_models
+ - Hardware Supported Models:
+ - models/hardware_supported_models/*
+ - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/
- Features: features
- Developer Guide:
- contributing/README.md
diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png
index f8c104ba14259..57a33524a5169 100644
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
diff --git a/docs/configuration/conserving_memory.md b/docs/configuration/conserving_memory.md
index 5ce43c7984057..0aa89a89eae5c 100644
--- a/docs/configuration/conserving_memory.md
+++ b/docs/configuration/conserving_memory.md
@@ -49,9 +49,6 @@ llm = LLM(model="adept/fuyu-8b", max_model_len=2048, max_num_seqs=2)
By default, we optimize model inference using CUDA graphs which take up extra memory in the GPU.
-!!! warning
- CUDA graph capture takes up more memory in V1 than in V0.
-
You can adjust `compilation_config` to achieve a better balance between inference speed and memory usage:
??? code
diff --git a/docs/configuration/env_vars.md b/docs/configuration/env_vars.md
index 2c0a898754fa0..f6d548a19d91f 100644
--- a/docs/configuration/env_vars.md
+++ b/docs/configuration/env_vars.md
@@ -7,8 +7,6 @@ vLLM uses the following environment variables to configure the system:
All environment variables used by vLLM are prefixed with `VLLM_`. **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables).
-??? code
-
- ```python
- --8<-- "vllm/envs.py:env-vars-definition"
- ```
+```python
+--8<-- "vllm/envs.py:env-vars-definition"
+```
diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md
index b0d390d7e1cbb..fdd9c317b022f 100644
--- a/docs/configuration/optimization.md
+++ b/docs/configuration/optimization.md
@@ -31,9 +31,7 @@ In vLLM V1, the default preemption mode is `RECOMPUTE` rather than `SWAP`, as re
Chunked prefill allows vLLM to process large prefills in smaller chunks and batch them together with decode requests. This feature helps improve both throughput and latency by better balancing compute-bound (prefill) and memory-bound (decode) operations.
-In vLLM V1, **chunked prefill is always enabled by default**. This is different from vLLM V0, where it was conditionally enabled based on model characteristics.
-
-With chunked prefill enabled, the scheduling policy prioritizes decode requests. It batches all pending decode requests before scheduling any prefill operations. When there are available tokens in the `max_num_batched_tokens` budget, it schedules pending prefills. If a pending prefill request cannot fit into `max_num_batched_tokens`, it automatically chunks it.
+In V1, **chunked prefill is enabled by default whenever possible**. With chunked prefill enabled, the scheduling policy prioritizes decode requests. It batches all pending decode requests before scheduling any prefill operations. When there are available tokens in the `max_num_batched_tokens` budget, it schedules pending prefills. If a pending prefill request cannot fit into `max_num_batched_tokens`, it automatically chunks it.
This policy has two benefits:
diff --git a/docs/configuration/tpu.md b/docs/configuration/tpu.md
deleted file mode 100644
index 2d24c9c6e2e95..0000000000000
--- a/docs/configuration/tpu.md
+++ /dev/null
@@ -1,111 +0,0 @@
-# TPU Optimization Tips
-
-This doc serves as a collection of handy tips for optimizing your vLLM on TPU workload.
-
-## Get started
-
-Looking for setup and installation instructions? Find them [here](https://docs.vllm.ai/projects/tpu/en/latest/getting_started/installation/).
-
-### TPU workload sizing
-
-When selecting the ideal number of chips for a single serving instance, it's important to account for both the model size and the average request context length. Adequate HBM for the KV cache is essential to ensure a sufficient number of concurrent requests can be processed.
-
-The following colab [calculator](https://colab.research.google.com/github/ericehanley/rightsize-vllm/blob/main/HBM_Calculator.ipynb) will tell you:
-
-- KV cache size requirement per token and per request
-- TPU/GPU memory consumed by the model weights
-- TPU/GPU memory allocated for the KV cache
-- Maximum \# of requests you can approximately set (--max-num-seqs)
-
-This approach serves as a general rule of thumb.
-
-#### Latency-throughput tradeoff
-
-As with rightsizing the number of chips for your workload, consider adjusting `--max-num-seqs` to fine-tune the latency-throughput balance. Decreasing `--max-num-seqs` and/or increasing the number of chips can help reduce latency.
-
-`--max-num-seqs` defines the number of concurrent decode slots, effectively limiting the number of requests the server can process tokens for simultaneously. Increasing this value allows the server to pre-allocate more HBM to handle a higher number of concurrent requests, which can maximize overall throughput. However, this often increases the end-to-end (e2e) latency per request.
-
-Therefore, carefully tuning `--max-num-seqs` is crucial to achieving the desired balance between latency and throughput for your specific workload.
-
-In a similar way, `--max-num-batch-tokens` can be adjusted down to improve latency, or adjusted up to improve throughput.
-
-#### Compilation and Caching
-
-Coming from a GPU background, one of the key differences you'll notice with TPUs is an initial compilation step. TPUs are specialized accelerators (ASICs) that achieve maximum performance by executing pre-compiled, static computation graphs via the XLA compiler. Unlike GPUs, which can handle dynamic input shapes more flexibly, TPUs require a specific compiled graph for each tensor shape (e.g., batch size and sequence length) they process.
-
-To manage this, vLLM performs a one-time "warmup" process when you first launch the server. During this phase, it pre-compiles the model for various common input shapes and saves these compiled graphs to a cache on disk or remote storage (located at `~/.cache/vllm/xla_cache` by default). This process can range significantly, anywhere from a few minutes to an hour depending on the size of the model and context length used.
-
-Although the first compilation can take some time, for all subsequent server launches, vLLM can load these graphs directly from the cache, eliminating the compilation time for future runs.
-
-Use `VLLM_XLA_CACHE_PATH` environment variable to write to shareable storage for future deployed nodes (like when using autoscaling).
-
-#### Reducing compilation time
-
-This initial compilation time ranges significantly and is impacted by many of the arguments discussed in this optimization doc. Factors that influence the length of time to compile are things like model size and `--max-num-batch-tokens`. Other arguments you can tune are things like `VLLM_TPU_MOST_MODEL_LEN`.
-
-### Optimize based on your data
-
-#### max-model-len vs. most-model-len
-
-
-
-If most of your requests are shorter than the maximum model length but you still need to accommodate occasional longer requests, setting a high maximum model length can negatively impact performance. In these cases, you can try introducing most-model-len by specifying the `VLLM_TPU_MOST_MODEL_LEN` environment variable.
-
-For example, 1% requests are 32k length and 99% requests are 2k length. You can pass 32k into `--max-model-len 32768` and use `VLLM_TPU_MOST_MODEL_LEN=2048`.
-
-The requests get subdivided into max-model-len and most-model-len categories, for the latter category, you can gain better performance since the server can process more requests at a time.
-
-#### Padding
-
-For online serving with latency requirements, consider switching to bucket padding by setting the `VLLM_TPU_BUCKET_PADDING_GAP` environment variable. Because of the layout of the TPU, try using increments of 128 (e.g., 128, 256, etc.)
-
-The server pads the requests into fixed lengths before sending them to the model to avoid recompilation. To read more about TPU padding, see [here](https://cloud.google.com/tpu/docs/performance-guide#xla-efficiencies). Currently, there are 2 ways to pad the requests:
-
-1. the default exponential padding (pad to the nearest power of 2)
-2. bucket padding (pad to the nearest linearly increasing bucket).
-
-When using bucket padding, the buckets start from 16, end at max_model_len, and increment by `VLLM_TPU_BUCKET_PADDING_GAP`.
-
-For example, max_model_len=512, padding_gap=64, the buckets will be [16, 32, 64, 128, 192, 256, 320, 384, 448, 512].
-
-The fewer tokens you pad, the less unnecessary computation TPU does, the better performance you can get. For example, if num_tokens=300, with exponential padding, you pad to 512, with the bucket_padding above, you pad to 320.
-
-However, you need to be careful to choose the padding gap. If the gap is too small, it means the number of buckets is large, leading to increased warmup (precompile) time and higher memory to store the compiled graph. Too many compiled graphs may lead to HBM OOM. Conversely, an overly large gap yields no performance improvement compared to the default exponential padding.
-
-#### Quantization
-
-If possible, use the precision that matches the chip’s hardware acceleration:
-
-- v5e has int4/int8 hardware acceleration in the MXU
-- v6e has int4/int8 hardware acceleration in the MXU
-
-Supported quantized formats and features in vLLM on TPU [Jul '25]:
-
-- INT8 W8A8
-- INT8 W8A16
-- FP8 KV cache
-- [WIP] FP8 W8A8
-- [WIP] AWQ
-- [WIP] FP4 W4A8
-
-#### Parallelization
-
-Don't set TP to be less than the number of chips on a single-host deployment.
-
-Although it’s common to do this with GPUs, don't try to fragment 2 or 8 different workloads across 8 chips on a single host. If you need 1 or 4 chips, just create an instance with 1 or 4 chips (these are partial-host machine types).
-
-### Tune your workloads
-
-Although we try to have great default configs, we strongly recommend you check out the [vLLM auto-tuner](../../benchmarks/auto_tune/README.md) to optimize your workloads for your use case.
-
-### Future Topics We'll Cover
-
-#### Profiling
-
-The auto-tuner provides a profile of optimized configurations as its final step. However, interpreting this profile can be challenging for new users. We plan to expand this section in the future with more detailed guidance. In the meantime, you can learn how to collect a TPU profile using vLLM's native profiling tools [here](../examples/offline_inference/profiling_tpu.md). This profile can provide valuable insights into your workload's performance.
-
-#### SPMD
-
-More details to come.
-
-**Want us to cover something that isn't listed here? Open up an issue please and cite this doc. We'd love to hear your questions or tips.**
diff --git a/docs/contributing/ci/update_pytorch_version.md b/docs/contributing/ci/update_pytorch_version.md
index 09fd85a466eed..735bb2e205332 100644
--- a/docs/contributing/ci/update_pytorch_version.md
+++ b/docs/contributing/ci/update_pytorch_version.md
@@ -98,21 +98,6 @@ to warm it up so that future builds are faster.
-## Update dependencies
-
-Several vLLM dependencies like xFormers depend on PyTorch and need
-to be updated accordingly. Rather than waiting for all of them to publish new
-releases (which would take too much time), they can be built from
-source to unblock the update process.
-
-### xFormers
-
-```bash
-export TORCH_CUDA_ARCH_LIST='7.5 8.0+PTX 9.0a'
-MAX_JOBS=16 uv pip install --system \
- --no-build-isolation "git+https://github.com/facebookresearch/xformers@v0.0.32.post2"
-```
-
## Update all the different vLLM platforms
Rather than attempting to update all vLLM platforms in a single pull request, it's more manageable
diff --git a/docs/contributing/model/basic.md b/docs/contributing/model/basic.md
index a7b54f015c2da..e828de0adf3c2 100644
--- a/docs/contributing/model/basic.md
+++ b/docs/contributing/model/basic.md
@@ -133,8 +133,6 @@ We consider 3 different scenarios:
For case (1), we recommend looking at the implementation of [`MambaForCausalLM`](../../../vllm/model_executor/models/mamba.py) (for Mamba-1) or [`Mamba2ForCausalLM`](../../../vllm/model_executor/models/mamba2.py) (for Mamba-2) as a reference.
The model should inherit protocol `IsAttentionFree` and also implement class methods `get_mamba_state_dtype_from_config` and `get_mamba_state_shape_from_config` to calculate the state shapes and data types from the config.
For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
-Please *do not* use the `MambaCacheManager` (deprecated in V1) or replicate any of the V0-specific code paths in the existing model implementations.
-V0-only classes and code will be removed in the very near future.
The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
@@ -146,6 +144,7 @@ We use "mamba-like" to refer to layers that posses a state that is updated in-pl
For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers.
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
+It is also worth noting that we should update `MAMBA_TYPE_TO_BACKEND_MAP` and `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/attention/backends/registry.py) when adding a new mamba backend.
Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.
diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md
index 7941b1f49ee8b..7634cc0859edf 100644
--- a/docs/contributing/profiling.md
+++ b/docs/contributing/profiling.md
@@ -224,6 +224,6 @@ snakeviz expensive_function.prof
Leverage VLLM_GC_DEBUG environment variable to debug GC costs.
-- VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elpased times
+- VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elapsed times
- VLLM_GC_DEBUG='{"top_objects":5}': enable GC debugger to log top 5
collected objects for each gc.collect
diff --git a/docs/deployment/docker.md b/docs/deployment/docker.md
index 1c639f3533d47..0e636c87f38a4 100644
--- a/docs/deployment/docker.md
+++ b/docs/deployment/docker.md
@@ -82,8 +82,7 @@ DOCKER_BUILDKIT=1 docker build . \
## Building for Arm64/aarch64
-A docker container can be built for aarch64 systems such as the Nvidia Grace-Hopper. At time of this writing, this requires the use
-of PyTorch Nightly and should be considered **experimental**. Using the flag `--platform "linux/arm64"` will attempt to build for arm64.
+A docker container can be built for aarch64 systems such as the Nvidia Grace-Hopper. At time of this writing, this should be considered **experimental**. Using the flag `--platform "linux/arm64"` will attempt to build for arm64.
!!! note
Multiple modules must be compiled, so this process can take a while. Recommend using `--build-arg max_jobs=` & `--build-arg nvcc_threads=`
@@ -94,7 +93,6 @@ of PyTorch Nightly and should be considered **experimental**. Using the flag `--
```bash
# Example of building on Nvidia GH200 server. (Memory usage: ~15GB, Build time: ~1475s / ~25 min, Image size: 6.93GB)
- python3 use_existing_torch.py
DOCKER_BUILDKIT=1 docker build . \
--file docker/Dockerfile \
--target vllm-openai \
@@ -102,7 +100,8 @@ of PyTorch Nightly and should be considered **experimental**. Using the flag `--
-t vllm/vllm-gh200-openai:latest \
--build-arg max_jobs=66 \
--build-arg nvcc_threads=2 \
- --build-arg torch_cuda_arch_list="9.0 10.0+PTX"
+ --build-arg torch_cuda_arch_list="9.0 10.0+PTX" \
+ --build-arg RUN_WHEEL_CHECK=false
```
!!! note
diff --git a/docs/deployment/frameworks/skypilot.md b/docs/deployment/frameworks/skypilot.md
index f4a984a6433e2..e9b0d5f0671c3 100644
--- a/docs/deployment/frameworks/skypilot.md
+++ b/docs/deployment/frameworks/skypilot.md
@@ -4,7 +4,7 @@
-vLLM can be **run and scaled to multiple service replicas on clouds and Kubernetes** with [SkyPilot](https://github.com/skypilot-org/skypilot), an open-source framework for running LLMs on any cloud. More examples for various open models, such as Llama-3, Mixtral, etc, can be found in [SkyPilot AI gallery](https://skypilot.readthedocs.io/en/latest/gallery/index.html).
+vLLM can be **run and scaled to multiple service replicas on clouds and Kubernetes** with [SkyPilot](https://github.com/skypilot-org/skypilot), an open-source framework for running LLMs on any cloud. More examples for various open models, such as Llama-3, Mixtral, etc., can be found in [SkyPilot AI gallery](https://skypilot.readthedocs.io/en/latest/gallery/index.html).
## Prerequisites
diff --git a/docs/design/debug_vllm_compile.md b/docs/design/debug_vllm_compile.md
index 3b454e851b54e..8912eb58f8ac7 100644
--- a/docs/design/debug_vllm_compile.md
+++ b/docs/design/debug_vllm_compile.md
@@ -9,7 +9,7 @@ TL;DR:
|----------|----------|-------------|
| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs |
| -O.mode=0 | mode=CompilationMode.NONE | Turn off torch.compile only |
-| -O.cudagraph_mode=NONE | compilation_config=CompilationConfig(mode=CompilationMode.NONE) | Turn off CUDAGraphs only |
+| -O.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only |
| -O.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor |
## vLLM-torch.compile overview
diff --git a/docs/design/io_processor_plugins.md b/docs/design/io_processor_plugins.md
index 2f4b17f191a5d..91ab4deae71df 100644
--- a/docs/design/io_processor_plugins.md
+++ b/docs/design/io_processor_plugins.md
@@ -1,6 +1,6 @@
# IO Processor Plugins
-IO Processor plugins are a feature that allows pre and post processing of the model input and output for pooling models. The idea is that users are allowed to pass a custom input to vLLM that is converted into one or more model prompts and fed to the model `encode` method. One potential use-case of such plugins is that of using vLLM for generating multi-modal data. Say users feed an image to vLLM and get an image in output.
+IO Processor plugins are a feature that allows pre- and post-processing of the model input and output for pooling models. The idea is that users are allowed to pass a custom input to vLLM that is converted into one or more model prompts and fed to the model `encode` method. One potential use-case of such plugins is that of using vLLM for generating multi-modal data. Say users feed an image to vLLM and get an image in output.
When performing an inference with IO Processor plugins, the prompt type is defined by the plugin and the same is valid for the final request output. vLLM does not perform any validation of input/output data, and it is up to the plugin to ensure the correct data is being fed to the model and returned to the user. As of now these plugins support only pooling models and can be triggered via the `encode` method in `LLM` and `AsyncLLM`, or in online serving mode via the `/pooling` endpoint.
diff --git a/docs/design/logits_processors.md b/docs/design/logits_processors.md
index acf7fc245462c..8eadeb386fcf2 100644
--- a/docs/design/logits_processors.md
+++ b/docs/design/logits_processors.md
@@ -411,7 +411,7 @@ Logits processor `update_state()` implementations should assume the following mo
* **"Condense" the batch to be contiguous:** starting with the lowest-index empty slot (which was caused by a Remove), apply a Unidirectional Move from the current highest non-empty slot in the batch to fill the empty slot. Proceed with additional Unidirectional Move operations in order of increasing empty slot destination index and decreasing non-empty slot source index until the batch is contiguous
- * **Shrink the batch:** a side-effect of condensing the batch is that empty slots resulting from Remove operations are grouped in a contiguous block at the end of the batch array. Thus, after condensing, update `BatchUpdate.batch_size` to reflect the number of non-empty slots
+ * **Shrink the batch:** a side effect of condensing the batch is that empty slots resulting from Remove operations are grouped in a contiguous block at the end of the batch array. Thus, after condensing, update `BatchUpdate.batch_size` to reflect the number of non-empty slots
5. Reorder the batch for improved efficiency. Depending on the attention backend implementation and the current characteristics of the batch, zero or more Swap Move operations may be applied to reorder the batch
@@ -548,7 +548,7 @@ Built-in logits processors are always loaded when the vLLM engine starts. See th
Review these logits processor implementations for guidance on writing built-in logits processors.
-Additionally, the following logits-processor-like functionalities are hard-coded into the sampler and do not yet utilize the programming model described above. Most of them will be refactored to use the aforemented logits processor programming model.
+Additionally, the following logits-processor-like functionalities are hard-coded into the sampler and do not yet utilize the programming model described above. Most of them will be refactored to use the aforementioned logits processor programming model.
* Allowed token IDs
diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md
index 7663b82266f0b..f0d5a3e934f39 100644
--- a/docs/design/moe_kernel_features.md
+++ b/docs/design/moe_kernel_features.md
@@ -1,22 +1,22 @@
-# Fused MoE Kernel features
+# Fused MoE Kernel Features
The purpose of this document is to provide an overview of the various MoE kernels (both modular and non-modular) so it will be easier to select an appropriate set of kernels for any particular situation. This includes information about the all2all backends used by modular kernels.
## Fused MoE Modular All2All backends
-There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalize` sub-classes provide an interface for each all2all backend.
+There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalize` subclasses provide an interface for each all2all backend.
The following table describes the relevant features of each backend, i.e. activation format, supported quantization schemes and async support.
-The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalize` subclass, the finalize step requires the same format. All the backend `prepare` methods expect activations in standard format and all the `finalize methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document.
+The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalize` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document.
-The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalize` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports. e.g. deepep_high_throughput supports only block-quantized fp8 format, any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 w/per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16.
+The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalize` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16.
Async backends support the use of DBO (Dual Batch Overlap) and shared expert overlap (where shared experts are computed during the combine step).
-Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalize` subclass, for non-modular kernels, it is up to the experts function to deal with this flag.
+Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalize` subclass. For non-modular kernels, it is up to the experts function to deal with this flag.
-unless otherwise specified, backends are controlled via `VLLM_ALL2ALL_BACKEND`. All backends except `flashinfer` only work with EP+DP or EP+TP. `Flashinfer` can work with EP or DP w/o EP.
+Unless otherwise specified, backends are controlled via `VLLM_ALL2ALL_BACKEND`. All backends except `flashinfer` only work with EP+DP or EP+TP. `Flashinfer` can work with EP or DP without EP.
-| Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Sub-class |
-|---------------------------------------|--------------------|-----------------|------------------------|-------|-----------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward_impl] |
-| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
-| deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
-| deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
-| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferAllToAllMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferAllToAllMoEPrepareAndFinalize] |
-| flashinfer4 | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferCutlassMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferCutlassMoEPrepareAndFinalize] |
-| flashinfer4 | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferCutlassMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferCutlassMoEPrepareAndFinalize] |
-| MoEPrepareAndFinalizeNoEP5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] |
-| BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] |
+| Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
+|---------|--------------------|--------------|---------------|-------|-----------------------|-----------|
+| naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward_impl] |
+| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
+| deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
+| deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
+| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferAllToAllMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferAllToAllMoEPrepareAndFinalize] |
+| flashinfer4 | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferCutlassMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferCutlassMoEPrepareAndFinalize] |
+| MoEPrepareAndFinalizeNoEP5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] |
+| BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
2. A,T quantization occurs after dispatch.
3. All quantization happens after dispatch.
4. Controlled by different env vars (`VLLM_FLASHINFER_MOE_BACKEND` "throughput" or "latency")
- 5. This is a no-op dispatcher that can be used to pair with any modular experts to produce a modular kernel that runs w/o dispatch or combine. These cannot be selected via environment variable. These are generally use for testing or adapting an expert subclass to the `fused_experts` API.
+ 5. This is a no-op dispatcher that can be used to pair with any modular experts to produce a modular kernel that runs without dispatch or combine. These cannot be selected via environment variable. These are generally use for testing or adapting an expert subclass to the `fused_experts` API.
6. This depends on the experts implementation.
---
@@ -66,44 +65,43 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
-## Fused MoE Experts Kernels
+## Fused Experts Kernels
-The are a number of MoE experts kernel implementations for different quantization types and architectures. Most follow the general API of the base Triton [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts] function. Many have modular kernel adapters so they can be used with compatible all2all backends. This table lists each experts kernel and its particular properties.
+There are a number of MoE experts kernel implementations for different quantization types and architectures. Most follow the general API of the base Triton [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts] function. Many have modular kernel adapters, so they can be used with compatible all2all backends. This table lists each experts kernel and its particular properties.
-Each kernel must be provided with one of the supported input activation formats. Some flavors of kernels support both standard and batched formats through different entry points, e.g. `TritonExperts` and `BatchedTritonExperts`. Batched format kernels are currently only needed for matching with certain all2all backends, e.g. `pplx`, `DeepEPLLPrepareAndFinalize`.
+Each kernel must be provided with one of the supported input activation formats. Some flavors of kernels support both standard and batched formats through different entry points, e.g. `TritonExperts` and `BatchedTritonExperts`. Batched format kernels are currently only needed for matching with certain all2all backends, e.g. `pplx` and `DeepEPLLPrepareAndFinalize`.
Similar to the backend kernels, each experts kernel only supports certain quantization formats. For non-modular experts, the activations will be in the original type and quantized internally by the kernel. Modular experts will expect the activations to already be in the quantized format. Both types of experts will yield outputs in the original activation type.
-Each experts kernel supports one or more activation functions, e.g. silu, gelu that are applied to the intermediate results.
+Each experts kernel supports one or more activation functions, e.g. silu or gelu, which are applied to the intermediate results.
As with the backends, some experts support applying topk weights on the input activations. The entries in the column in this table only apply to the non-modular experts.
Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEPermuteExpertsUnpermute`.
-To be used with a particular `FusedMoEPrepareAndFinalize` sub-class, MoE kernels must have compatible activation formats, quantization types and quantization formats.
+To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats.
-| Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
-|------------------------------|-----------------------|------------------|---------------|-------------------------------------------------------------|-----------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| triton | standard | all1 | G,A,T | silu, gelu,swigluoai,silu_no_mul,gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
-| triton (batched) | batched | all1 | G,A,T | silu, gelu | 6 | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
-| deep gemm | standard,batched | fp8 | G(128),A,T | silu, gelu | 6 | Y | [`deep_gemm_moe_fp8`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.deep_gemm_moe_fp8],[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
-| cutlass_fp4 | standard,batched | nvfp4 | A,T | silu | Y | Y | [`cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp4],[`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
-| cutlass_fp8 | standard,batched | fp8 | A,T | silu, gelu | Y | Y | [`cutlass_moe_fp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp8],[`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
-| flashinfer | standard | nvfp4,fp8 | T | 5 | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
-| gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
-| deep gemm+triton2 | standard,batched | all1 | G(128),A,T | silu, gelu | 6 | Y | [`TritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe.TritonOrDeepGemmExperts],[`BatchedTritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe.BatchedTritonOrDeepGemmExperts] |
-| marlin | standard | 3 | 3 | silu,swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
-| marlin experts | standard,batched | N/A | N/A | silu,swigluoai | Y | Y | [`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
-| trtllm | standard | mxfp4,nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
-| pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
-| iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
-| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
-| 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.fused_batched_moe.NaiveBatchedExperts] |
+| Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
+|--------|-------------------|--------------|---------------|---------------------|-----------------------|---------|--------|
+| triton | standard | all1 | G,A,T | silu, gelu,swigluoai,silu_no_mul,gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
+| triton (batched) | batched | all1 | G,A,T | silu, gelu | 6 | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
+| deep gemm | standard,batched | fp8 | G(128),A,T | silu, gelu | 6 | Y | [`deep_gemm_moe_fp8`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.deep_gemm_moe_fp8],[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
+| cutlass_fp4 | standard,batched | nvfp4 | A,T | silu | Y | Y | [`cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp4],[`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
+| cutlass_fp8 | standard,batched | fp8 | A,T | silu, gelu | Y | Y | [`cutlass_moe_fp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp8],[`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
+| flashinfer | standard | nvfp4,fp8 | T | 5 | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
+| gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
+| deep gemm+triton2 | standard,batched | all1 | G(128),A,T | silu, gelu | 6 | Y | [`TritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe.TritonOrDeepGemmExperts],[`BatchedTritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe.BatchedTritonOrDeepGemmExperts] |
+| marlin | standard,batched | 3 / N/A | 3 / N/A | silu,swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
+| trtllm | standard | mxfp4,nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
+| pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
+| iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
+| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
+| 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.fused_batched_moe.NaiveBatchedExperts] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
- 2. A dispatcher wrapper around triton and deep gemm experts. Will select based on type + shape + quantization params
+ 2. A dispatcher wrapper around triton and deep gemm experts. Will select based on type + shape + quantization params
3. uint4, uint8, fp8, fp4
4. This is a naive implementation of experts that supports batched format. Mainly used for testing.
5. The `activation` parameter is ignored and SwiGlu is used by default instead.
@@ -113,8 +111,8 @@ To be used with a particular `FusedMoEPrepareAndFinalize` sub-class, MoE kernels
The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts.
-| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
-|----------------------------------|------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------|
-| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,`TritonExperts`,`TritonOrDeepGemmExperts`,`CutlassExpertsFp8`, `MarlinExperts` |
-| deepep_low_latency,pplx | `DeepEPLLPrepareAndFinalize`,`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,`BatchedTritonExperts`,`BatchedTritonOrDeepGemmExperts`,`CutlassBatchedExpertsFp8`,`BatchedMarlinExperts`|
-| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
+| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
+|---------|-----------------------------------------|----------------------------------------------|
+| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,`TritonExperts`,`TritonOrDeepGemmExperts`,`CutlassExpertsFp8`, `MarlinExperts` |
+| deepep_low_latency,pplx | `DeepEPLLPrepareAndFinalize`,`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,`BatchedTritonExperts`,`BatchedTritonOrDeepGemmExperts`,`CutlassBatchedExpertsFp8`,`BatchedMarlinExperts` |
+| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
diff --git a/docs/design/plugin_system.md b/docs/design/plugin_system.md
index dc2f7c4aed3c3..9c84889f7f03d 100644
--- a/docs/design/plugin_system.md
+++ b/docs/design/plugin_system.md
@@ -4,7 +4,7 @@ The community frequently requests the ability to extend vLLM with custom feature
## How Plugins Work in vLLM
-Plugins are user-registered code that vLLM executes. Given vLLM's architecture (see [Arch Overview](arch_overview.md)), multiple processes may be involved, especially when using distributed inference with various parallelism techniques. To enable plugins successfully, every process created by vLLM needs to load the plugin. This is done by the [load_general_plugins](https://github.com/vllm-project/vllm/blob/c76ac49d266e27aa3fea84ef2df1f813d24c91c7/vllm/plugins/__init__.py#L16) function in the `vllm.plugins` module. This function is called for every process created by vLLM before it starts any work.
+Plugins are user-registered code that vLLM executes. Given vLLM's architecture (see [Arch Overview](arch_overview.md)), multiple processes may be involved, especially when using distributed inference with various parallelism techniques. To enable plugins successfully, every process created by vLLM needs to load the plugin. This is done by the [load_plugins_by_group][vllm.plugins.load_plugins_by_group] function in the `vllm.plugins` module.
## How vLLM Discovers Plugins
@@ -49,7 +49,7 @@ Every plugin has three parts:
- **Platform plugins** (with group name `vllm.platform_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree platforms into vLLM. The plugin function should return `None` when the platform is not supported in the current environment, or the platform class's fully qualified name when the platform is supported.
-- **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre/post processing of the model prompt and model output for pooling models. The plugin function returns the IOProcessor's class fully qualified name.
+- **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre-/post-processing of the model prompt and model output for pooling models. The plugin function returns the IOProcessor's class fully qualified name.
- **Stat logger plugins** (with group name `vllm.stat_logger_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree loggers into vLLM. The entry point should be a class that subclasses StatLoggerBase.
@@ -57,6 +57,100 @@ Every plugin has three parts:
- **Being re-entrant**: The function specified in the entry point should be re-entrant, meaning it can be called multiple times without causing issues. This is necessary because the function might be called multiple times in some processes.
+### Platform plugins guidelines
+
+1. Create a platform plugin project, for example, `vllm_add_dummy_platform`. The project structure should look like this:
+
+ ```shell
+ vllm_add_dummy_platform/
+ ├── vllm_add_dummy_platform/
+ │ ├── __init__.py
+ │ ├── my_dummy_platform.py
+ │ ├── my_dummy_worker.py
+ │ ├── my_dummy_attention.py
+ │ ├── my_dummy_device_communicator.py
+ │ ├── my_dummy_custom_ops.py
+ ├── setup.py
+ ```
+
+2. In the `setup.py` file, add the following entry point:
+
+ ```python
+ setup(
+ name="vllm_add_dummy_platform",
+ ...
+ entry_points={
+ "vllm.platform_plugins": [
+ "my_dummy_platform = vllm_add_dummy_platform:register"
+ ]
+ },
+ ...
+ )
+ ```
+
+ Please make sure `vllm_add_dummy_platform:register` is a callable function and returns the platform class's fully qualified name. for example:
+
+ ```python
+ def register():
+ return "vllm_add_dummy_platform.my_dummy_platform.MyDummyPlatform"
+ ```
+
+3. Implement the platform class `MyDummyPlatform` in `my_dummy_platform.py`. The platform class should inherit from `vllm.platforms.interface.Platform`. Please follow the interface to implement the functions one by one. There are some important functions and properties that should be implemented at least:
+
+ - `_enum`: This property is the device enumeration from [PlatformEnum][vllm.platforms.interface.PlatformEnum]. Usually, it should be `PlatformEnum.OOT`, which means the platform is out-of-tree.
+ - `device_type`: This property should return the type of the device which pytorch uses. For example, `"cpu"`, `"cuda"`, etc.
+ - `device_name`: This property is set the same as `device_type` usually. It's mainly used for logging purposes.
+ - `check_and_update_config`: This function is called very early in the vLLM's initialization process. It's used for plugins to update the vllm configuration. For example, the block size, graph mode config, etc, can be updated in this function. The most important thing is that the **worker_cls** should be set in this function to let vLLM know which worker class to use for the worker process.
+ - `get_attn_backend_cls`: This function should return the attention backend class's fully qualified name.
+ - `get_device_communicator_cls`: This function should return the device communicator class's fully qualified name.
+
+4. Implement the worker class `MyDummyWorker` in `my_dummy_worker.py`. The worker class should inherit from [WorkerBase][vllm.v1.worker.worker_base.WorkerBase]. Please follow the interface to implement the functions one by one. Basically, all interfaces in the base class should be implemented, since they are called here and there in vLLM. To make sure a model can be executed, the basic functions should be implemented are:
+
+ - `init_device`: This function is called to set up the device for the worker.
+ - `initialize_cache`: This function is called to set cache config for the worker.
+ - `load_model`: This function is called to load the model weights to device.
+ - `get_kv_cache_spaces`: This function is called to generate the kv cache spaces for the model.
+ - `determine_available_memory`: This function is called to profiles the peak memory usage of the model to determine how much memory can be used for KV cache without OOMs.
+ - `initialize_from_config`: This function is called to allocate device KV cache with the specified kv_cache_config
+ - `execute_model`: This function is called every step to inference the model.
+
+ Additional functions that can be implemented are:
+
+ - If the plugin wants to support sleep mode feature, please implement the `sleep` and `wakeup` functions.
+ - If the plugin wants to support graph mode feature, please implement the `compile_or_warm_up_model` function.
+ - If the plugin wants to support speculative decoding feature, please implement the `take_draft_token_ids` function.
+ - If the plugin wants to support lora feature, please implement the `add_lora`,`remove_lora`,`list_loras` and `pin_lora` functions.
+ - If the plugin wants to support data parallelism feature, please implement the `execute_dummy_batch` functions.
+
+ Please look at the worker base class [WorkerBase][vllm.v1.worker.worker_base.WorkerBase] for more functions that can be implemented.
+
+5. Implement the attention backend class `MyDummyAttention` in `my_dummy_attention.py`. The attention backend class should inherit from [AttentionBackend][vllm.attention.backends.abstract.AttentionBackend]. It's used to calculate attentions with your device. Take `vllm.v1.attention.backends` as examples, it contains many attention backend implementations.
+
+6. Implement custom ops for high performance. Most ops can be ran by pytorch native implementation, while the performance may not be good. In this case, you can implement specific custom ops for your plugins. Currently, there are kinds of custom ops vLLM supports:
+
+ - pytorch ops
+ there are 3 kinds of pytorch ops:
+
+ - `communicator ops`: Device communicator op. Such as all-reduce, all-gather, etc.
+ Please implement the device communicator class `MyDummyDeviceCommunicator` in `my_dummy_device_communicator.py`. The device communicator class should inherit from [DeviceCommunicatorBase][vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase].
+ - `common ops`: Common ops. Such as matmul, softmax, etc.
+ Please implement the common ops by register oot way. See more detail in [CustomOp][vllm.model_executor.custom_op.CustomOp] class.
+ - `csrc ops`: C++ ops. This kind of ops are implemented in C++ and are registered as torch custom ops.
+ Following csrc module and `vllm._custom_ops` to implement your ops.
+
+ - triton ops
+ Custom way doesn't work for triton ops now.
+
+7. (optional) Implement other plugable modules, such as lora, graph backend, quantization, mamba attention backend, etc.
+
## Compatibility Guarantee
-vLLM guarantees the interface of documented plugins, such as `ModelRegistry.register_model`, will always be available for plugins to register models. However, it is the responsibility of plugin developers to ensure their plugins are compatible with the version of vLLM they are targeting. For example, `"vllm_add_dummy_model.my_llava:MyLlava"` should be compatible with the version of vLLM that the plugin targets. The interface for the model may change during vLLM's development.
+vLLM guarantees the interface of documented plugins, such as `ModelRegistry.register_model`, will always be available for plugins to register models. However, it is the responsibility of plugin developers to ensure their plugins are compatible with the version of vLLM they are targeting. For example, `"vllm_add_dummy_model.my_llava:MyLlava"` should be compatible with the version of vLLM that the plugin targets.
+
+The interface for the model/module may change during vLLM's development. If you see any deprecation log info, please upgrade your plugin to the latest version.
+
+## Deprecation announcement
+
+!!! warning "Deprecations"
+ - `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It will be removed in v0.13.0 or v1.0.0.
+ - `_Backend` in `vllm.attention` is deprecated. It will be removed in v0.13.0 or v1.0.0. Please use `vllm.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
diff --git a/docs/design/prefix_caching.md b/docs/design/prefix_caching.md
index bd4070f381d81..cf792fdabe1a6 100644
--- a/docs/design/prefix_caching.md
+++ b/docs/design/prefix_caching.md
@@ -1,6 +1,6 @@
# Automatic Prefix Caching
-Prefix caching kv-cache blocks is a popular optimization in LLM inference to avoid redundant prompt computations. The core idea is simple – we cache the kv-cache blocks of processed requests, and reuse these blocks when a new request comes in with the same prefix as previous requests. Since prefix caching is almost a free lunch and won’t change model outputs, it has been widely used by many public endpoints (e.g., OpenAI, Anthropic, etc) and most open source LLM inference frameworks (e.g., SGLang).
+Prefix caching kv-cache blocks is a popular optimization in LLM inference to avoid redundant prompt computations. The core idea is simple – we cache the kv-cache blocks of processed requests, and reuse these blocks when a new request comes in with the same prefix as previous requests. Since prefix caching is almost a free lunch and won’t change model outputs, it has been widely used by many public endpoints (e.g., OpenAI, Anthropic, etc.) and most open source LLM inference frameworks (e.g., SGLang).
While there are many ways to implement prefix caching, vLLM chooses a hash-based approach. Specifically, we hash each kv-cache block by the tokens in the block and the tokens in the prefix before the block:
@@ -94,9 +94,6 @@ To improve privacy in shared environments, vLLM supports isolating prefix cache
With this setup, cache sharing is limited to users or requests that explicitly agree on a common salt, enabling cache reuse within a trust group while isolating others.
-!!! note
- Cache isolation is not supported in engine V0.
-
## Data Structure
The prefix caching in vLLM v1 is implemented in the KV cache manager. The basic building block is the “Block” data class (simplified):
diff --git a/docs/features/README.md b/docs/features/README.md
index ad9de9ff8f368..5faf3768f3214 100644
--- a/docs/features/README.md
+++ b/docs/features/README.md
@@ -59,20 +59,23 @@ th:not(:first-child) {
### Feature x Hardware
-| Feature | Volta | Turing | Ampere | Ada | Hopper | CPU | AMD | TPU | Intel GPU |
-|-----------------------------------------------------------|---------------------|-----------|-----------|--------|------------|--------------------|--------|-----| ------------|
-| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| [SD](spec_decode.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | [🟠](https://github.com/vllm-project/vllm/issues/26963) |
-| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
-| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
-| enc-dec | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
-| [mm](multimodal_inputs.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | [🟠](https://github.com/vllm-project/vllm/issues/26965) |
-| logP | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
-| prmpt logP | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
-| async output | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ |
-| multi-step | ✅ | ✅ | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/8477) | ✅ | ❌ | ✅ |
-| best-of | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
-| beam-search | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
-| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/25097) | ✅ |
+| Feature | Volta | Turing | Ampere | Ada | Hopper | CPU | AMD | Intel GPU |
+|-----------------------------------------------------------|---------------------|-----------|-----------|--------|------------|--------------------|--------| ------------|
+| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| [SD](spec_decode.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [🟠](https://github.com/vllm-project/vllm/issues/26963) |
+| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
+| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| enc-dec | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
+| [mm](multimodal_inputs.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | [🟠](https://github.com/vllm-project/vllm/issues/26965) |
+| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ |
+| logP | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| prmpt logP | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| async output | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
+| multi-step | ✅ | ✅ | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/8477) | ✅ | ✅ |
+| best-of | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| beam-search | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+
+!!! note
+ For information on feature support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
diff --git a/docs/features/disagg_prefill.md b/docs/features/disagg_prefill.md
index 3e8cb87e37d33..fd4f249f2ec6c 100644
--- a/docs/features/disagg_prefill.md
+++ b/docs/features/disagg_prefill.md
@@ -91,6 +91,6 @@ Disaggregated prefilling is highly related to infrastructure, so vLLM relies on
We recommend three ways of implementations:
-- **Fully-customized connector**: Implement your own `Connector`, and call third-party libraries to send and receive KV caches, and many many more (like editing vLLM's model input to perform customized prefilling, etc). This approach gives you the most control, but at the risk of being incompatible with future vLLM versions.
+- **Fully-customized connector**: Implement your own `Connector`, and call third-party libraries to send and receive KV caches, and many many more (like editing vLLM's model input to perform customized prefilling, etc.). This approach gives you the most control, but at the risk of being incompatible with future vLLM versions.
- **Database-like connector**: Implement your own `LookupBuffer` and support the `insert` and `drop_select` APIs just like SQL.
- **Distributed P2P connector**: Implement your own `Pipe` and support the `send_tensor` and `recv_tensor` APIs, just like `torch.distributed`.
diff --git a/docs/features/lora.md b/docs/features/lora.md
index 3a85b52d89b68..d42a3cef76bde 100644
--- a/docs/features/lora.md
+++ b/docs/features/lora.md
@@ -4,7 +4,7 @@ This document shows you how to use [LoRA adapters](https://arxiv.org/abs/2106.09
LoRA adapters can be used with any vLLM model that implements [SupportsLoRA][vllm.model_executor.models.interfaces.SupportsLoRA].
-Adapters can be efficiently served on a per request basis with minimal overhead. First we download the adapter(s) and save
+Adapters can be efficiently served on a per-request basis with minimal overhead. First we download the adapter(s) and save
them locally with
```python
diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md
index cde2ec165712b..4656ee43ea251 100644
--- a/docs/features/multimodal_inputs.md
+++ b/docs/features/multimodal_inputs.md
@@ -365,6 +365,8 @@ You must enable this feature via `enable_mm_embeds=True`.
The vLLM engine may crash if incorrect shape of embeddings is passed.
Only enable this flag for trusted users!
+#### Image Embeddings
+
??? code
```python
@@ -441,6 +443,36 @@ For Qwen2-VL and MiniCPM-V, we accept additional parameters alongside the embedd
print(generated_text)
```
+#### Audio Embeddings
+
+You can pass pre-computed audio embeddings similar to image embeddings:
+
+??? code
+
+ ```python
+ from vllm import LLM
+ import torch
+
+ # Enable audio embeddings support
+ llm = LLM(model="fixie-ai/ultravox-v0_5-llama-3_2-1b", enable_mm_embeds=True)
+
+ # Refer to the HuggingFace repo for the correct format to use
+ prompt = "USER: \nWhat is in this audio?\nASSISTANT:"
+
+ # Load pre-computed audio embeddings
+ # torch.Tensor of shape (1, audio_feature_size, hidden_size of LM)
+ audio_embeds = torch.load(...)
+
+ outputs = llm.generate({
+ "prompt": prompt,
+ "multi_modal_data": {"audio": audio_embeds},
+ })
+
+ for o in outputs:
+ generated_text = o.outputs[0].text
+ print(generated_text)
+ ```
+
## Online Serving
Our OpenAI-compatible server accepts multi-modal data via the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat). Media inputs also support optional UUIDs users can provide to uniquely identify each media, which is used to cache the media results across requests.
@@ -483,7 +515,7 @@ Then, you can use the OpenAI client as follows:
)
# Single-image input inference
- image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
chat_response = client.chat.completions.create(
model="microsoft/Phi-3.5-vision-instruct",
diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md
index 1ce038f4d6525..f0e25e31aa0b3 100644
--- a/docs/features/nixl_connector_usage.md
+++ b/docs/features/nixl_connector_usage.md
@@ -158,7 +158,7 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
## Experimental Feature
-### Heterogenuous KV Layout support
+### Heterogeneous KV Layout support
Support use case: Prefill with 'HND' and decode with 'NHD' with experimental configuration
diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md
index 74f005c496ee5..7b5287bad3bb8 100644
--- a/docs/features/quantization/README.md
+++ b/docs/features/quantization/README.md
@@ -43,24 +43,27 @@ th:not(:first-child) {
}
-| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | Intel Gaudi | x86 CPU | Google TPU |
-|-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-------------|-----------|--------------|
-| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ | ❌ |
-| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ | ❌ |
-| Marlin (GPTQ/AWQ/FP8) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ✅︎ | ✅︎ |
-| FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
-| BitBLAS | ✅︎ | ✅ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| BitBLAS (GPTQ) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
-| INC (W8A8) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅︎ | ❌ | ❌ |
+| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | Intel Gaudi | x86 CPU |
+|-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-------------|-----------|
+| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ |
+| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ |
+| Marlin (GPTQ/AWQ/FP8) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
+| INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ✅︎ |
+| FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
+| BitBLAS | ✅︎ | ✅ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
+| BitBLAS (GPTQ) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
+| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
+| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
+| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
+| INC (W8A8) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅︎ | ❌ |
- Volta refers to SM 7.0, Turing to SM 7.5, Ampere to SM 8.0/8.6, Ada to SM 8.9, and Hopper to SM 9.0.
- ✅︎ indicates that the quantization method is supported on the specified hardware.
- ❌ indicates that the quantization method is not supported on the specified hardware.
+!!! note
+ For information on quantization support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
+
!!! note
This compatibility chart is subject to change as vLLM continues to evolve and expand its support for different hardware platforms and quantization methods.
diff --git a/docs/features/quantization/fp8.md b/docs/features/quantization/fp8.md
index 0c5111fb8af0d..d4a6176b236f1 100644
--- a/docs/features/quantization/fp8.md
+++ b/docs/features/quantization/fp8.md
@@ -60,7 +60,7 @@ Since simple RTN does not require data for weight quantization and the activatio
??? code
```python
- from llmcompressor.transformers import oneshot
+ from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
# Configure the simple PTQ quantization
diff --git a/docs/features/quantization/int4.md b/docs/features/quantization/int4.md
index 035e7ea291f9e..9752039097d63 100644
--- a/docs/features/quantization/int4.md
+++ b/docs/features/quantization/int4.md
@@ -80,7 +80,7 @@ Now, apply the quantization algorithms:
??? code
```python
- from llmcompressor.transformers import oneshot
+ from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
diff --git a/docs/features/quantization/int8.md b/docs/features/quantization/int8.md
index ec8a77f74ffef..701ca6378cb16 100644
--- a/docs/features/quantization/int8.md
+++ b/docs/features/quantization/int8.md
@@ -87,7 +87,7 @@ Now, apply the quantization algorithms:
??? code
```python
- from llmcompressor.transformers import oneshot
+ from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
diff --git a/docs/features/quantization/quantized_kvcache.md b/docs/features/quantization/quantized_kvcache.md
index 56cf057678be6..d26a5e217f314 100644
--- a/docs/features/quantization/quantized_kvcache.md
+++ b/docs/features/quantization/quantized_kvcache.md
@@ -78,7 +78,7 @@ Here's a complete example using `meta-llama/Llama-3.1-8B-Instruct` (most models
```python
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
- from llmcompressor.transformers import oneshot
+ from llmcompressor import oneshot
# Select model and load it
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
diff --git a/docs/features/quantization/quark.md b/docs/features/quantization/quark.md
index bd7bc186e13aa..c54d7d2251999 100644
--- a/docs/features/quantization/quark.md
+++ b/docs/features/quantization/quark.md
@@ -306,7 +306,7 @@ As examples, we provide some ready-to-use quantized mixed precision model to sho
### 2. inference the quantized mixed precision model in vLLM
-Models quantized with AMD Quark using mixed precision can natively be reload in vLLM, and e.g. evaluated using lm-evaluation-harness as follow:
+Models quantized with AMD Quark using mixed precision can natively be reload in vLLM, and e.g. evaluated using lm-evaluation-harness as follows:
```bash
lm_eval --model vllm \
diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md
index 7e6c69e717dba..dd79ba19b7247 100644
--- a/docs/features/tool_calling.md
+++ b/docs/features/tool_calling.md
@@ -142,7 +142,7 @@ Flags: `--tool-call-parser hermes`
Supported models:
* `mistralai/Mistral-7B-Instruct-v0.3` (confirmed)
-* Additional mistral function-calling models are compatible as well.
+* Additional Mistral function-calling models are compatible as well.
Known issues:
@@ -158,12 +158,25 @@ Known issues:
Recommended flags:
-1. To use [mistral-common](https://github.com/mistralai/mistral-common) the official Mistral tokenization backend:
+1. To use the official Mistral AI's format:
- `--tokenizer_mode mistral --config_format mistral --load_format mistral --tool-call-parser mistral`
+ `--tool-call-parser mistral`
-2. To use the default Transformers tokenization backend:
- `--tool-call-parser mistral --chat-template examples/tool_chat_template_mistral_parallel.jinja`
+2. To use the Transformers format when available:
+
+ `--tokenizer_mode hf --config_format hf --load_format hf --tool-call-parser mistral --chat-template examples/tool_chat_template_mistral_parallel.jinja`
+
+!!! note
+ Models officially released by Mistral AI have two possible formats:
+
+ 1. The official format that is used by default with `auto` or `mistral` arguments:
+
+ `--tokenizer_mode mistral --config_format mistral --load_format mistral`
+ This format uses [mistral-common](https://github.com/mistralai/mistral-common), the Mistral AI's tokenizer backend.
+
+ 2. The Transformers format, when available, that is used with `hf` arguments:
+
+ `--tokenizer_mode hf --config_format hf --load_format hf --chat-template examples/tool_chat_template_mistral_parallel.jinja`
### Llama Models (`llama3_json`)
diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md
index be99cef3723e6..d1beab7855b18 100644
--- a/docs/getting_started/installation/cpu.md
+++ b/docs/getting_started/installation/cpu.md
@@ -97,7 +97,6 @@ Currently, there are no pre-built CPU wheels.
- `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads, can be set as CPU id lists, `auto` (by default), or `nobind` (to disable binding to individual CPU cores and to inherit user-defined OpenMP variables). For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node respectively. If set to `nobind`, the number of OpenMP threads is determined by the standard `OMP_NUM_THREADS` environment variable.
- `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`.
- `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence.
-- `VLLM_CPU_MOE_PREPACK` (x86 only): whether to use prepack for MoE layer. This will be passed to `ipex.llm.modules.GatedMLPMOE`. Default is `1` (True). On unsupported CPUs, you might need to set this to `0` (False).
- `VLLM_CPU_SGL_KERNEL` (x86 only, Experimental): whether to use small-batch optimized kernels for linear layer and MoE layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False).
## FAQ
@@ -191,10 +190,9 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
- GPTQ (x86 only)
- compressed-tensor INT8 W8A8 (x86, s390x)
-### (x86 only) What is the purpose of `VLLM_CPU_MOE_PREPACK` and `VLLM_CPU_SGL_KERNEL`?
+### (x86 only) What is the purpose of `VLLM_CPU_SGL_KERNEL`?
- Both of them require `amx` CPU flag.
- - `VLLM_CPU_MOE_PREPACK` can provide better performance for MoE models
- `VLLM_CPU_SGL_KERNEL` can provide better performance for MoE models and small-batch scenarios.
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md
index b2d0d64a2d355..601d3659af886 100644
--- a/docs/getting_started/installation/gpu.cuda.inc.md
+++ b/docs/getting_started/installation/gpu.cuda.inc.md
@@ -158,10 +158,7 @@ uv pip install -e .
##### Use an existing PyTorch installation
-There are scenarios where the PyTorch dependency cannot be easily installed with `uv`, e.g.:
-
-- Building vLLM with PyTorch nightly or a custom PyTorch build.
-- Building vLLM with aarch64 and CUDA (GH200), where the PyTorch wheels are not available on PyPI. Currently, only the PyTorch nightly has wheels for aarch64 with CUDA. You can run `uv pip install --index-url https://download.pytorch.org/whl/nightly/cu128 torch torchvision torchaudio` to [install PyTorch nightly](https://pytorch.org/get-started/locally/) and then build vLLM on top of it.
+There are scenarios where the PyTorch dependency cannot be easily installed with `uv`, for example, when building vLLM with non-default PyTorch builds (like nightly or a custom build).
To build vLLM using an existing PyTorch installation:
diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md
index cfc8b4d9838a7..94920dc5306b3 100644
--- a/docs/getting_started/quickstart.md
+++ b/docs/getting_started/quickstart.md
@@ -283,10 +283,10 @@ Currently, vLLM supports multiple backends for efficient Attention computation a
If desired, you can also manually set the backend of your choice by configuring the environment variable `VLLM_ATTENTION_BACKEND` to one of the following options:
-- On NVIDIA CUDA: `FLASH_ATTN`, `FLASHINFER` or `XFORMERS`.
+- On NVIDIA CUDA: `FLASH_ATTN` or `FLASHINFER`.
- On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA` or `ROCM_AITER_UNIFIED_ATTN`.
-For AMD ROCm, you can futher control the specific Attention implementation using the following variables:
+For AMD ROCm, you can further control the specific Attention implementation using the following variables:
- Triton Unified Attention: `VLLM_ROCM_USE_AITER=0 VLLM_V1_USE_PREFILL_DECODE_ATTENTION=0 VLLM_ROCM_USE_AITER_MHA=0`
- AITER Unified Attention: `VLLM_ROCM_USE_AITER=1 VLLM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_V1_USE_PREFILL_DECODE_ATTENTION=0 VLLM_ROCM_USE_AITER_MHA=0`
diff --git a/docs/models/hardware_supported_models/cpu.md b/docs/models/hardware_supported_models/cpu.md
new file mode 100644
index 0000000000000..0832755f8fbe2
--- /dev/null
+++ b/docs/models/hardware_supported_models/cpu.md
@@ -0,0 +1,26 @@
+# CPU - Intel® Xeon®
+
+## Supported Models
+
+### Text-only Language Models
+
+| Model | Architecture | Supported |
+|--------------------------------------|-------------------------------------------|-----------|
+| meta-llama/Llama-3.1 / 3.3 | LlamaForCausalLM | ✅ |
+| meta-llama/Llama-4-Scout | Llama4ForConditionalGeneration | ✅ |
+| meta-llama/Llama-4-Maverick | Llama4ForConditionalGeneration | ✅ |
+| ibm-granite/granite (Granite-MOE) | GraniteMoeForCausalLM | ✅ |
+| Qwen/Qwen3 | Qwen3ForCausalLM | ✅ |
+| zai-org/GLM-4.5 | GLMForCausalLM | ✅ |
+| google/gemma | GemmaForCausalLM | ✅ |
+
+### Multimodal Language Models
+
+| Model | Architecture | Supported |
+|--------------------------------------|-------------------------------------------|-----------|
+| Qwen/Qwen2.5-VL | Qwen2VLForConditionalGeneration | ✅ |
+| openai/whisper | WhisperForConditionalGeneration | ✅ |
+
+✅ Runs and optimized.
+🟨 Runs and correct but not optimized to green yet.
+❌ Does not pass accuracy test or does not run.
diff --git a/docs/models/hardware_supported_models/tpu.md b/docs/models/hardware_supported_models/tpu.md
deleted file mode 100644
index 7b0a5ba6e72da..0000000000000
--- a/docs/models/hardware_supported_models/tpu.md
+++ /dev/null
@@ -1,34 +0,0 @@
-# TPU
-
-## Supported Models
-
-### Text-only Language Models
-
-| Model | Architecture | Supported |
-|-----------------------------------------------------|--------------------------------|-----------|
-| mistralai/Mixtral-8x7B-Instruct-v0.1 | MixtralForCausalLM | 🟨 |
-| mistralai/Mistral-Small-24B-Instruct-2501 | MistralForCausalLM | ✅ |
-| mistralai/Codestral-22B-v0.1 | MistralForCausalLM | ✅ |
-| mistralai/Mixtral-8x22B-Instruct-v0.1 | MixtralForCausalLM | ❌ |
-| meta-llama/Llama-3.3-70B-Instruct | LlamaForCausalLM | ✅ |
-| meta-llama/Llama-3.1-8B-Instruct | LlamaForCausalLM | ✅ |
-| meta-llama/Llama-3.1-70B-Instruct | LlamaForCausalLM | ✅ |
-| meta-llama/Llama-4-* | Llama4ForConditionalGeneration | ❌ |
-| microsoft/Phi-3-mini-128k-instruct | Phi3ForCausalLM | 🟨 |
-| microsoft/phi-4 | Phi3ForCausalLM | ❌ |
-| google/gemma-3-27b-it | Gemma3ForConditionalGeneration | 🟨 |
-| google/gemma-3-4b-it | Gemma3ForConditionalGeneration | ❌ |
-| deepseek-ai/DeepSeek-R1 | DeepseekV3ForCausalLM | ❌ |
-| deepseek-ai/DeepSeek-V3 | DeepseekV3ForCausalLM | ❌ |
-| RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
-| RedHatAI/Meta-Llama-3.1-70B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
-| Qwen/Qwen3-8B | Qwen3ForCausalLM | ✅ |
-| Qwen/Qwen3-32B | Qwen3ForCausalLM | ✅ |
-| Qwen/Qwen2.5-7B-Instruct | Qwen2ForCausalLM | ✅ |
-| Qwen/Qwen2.5-32B | Qwen2ForCausalLM | ✅ |
-| Qwen/Qwen2.5-14B-Instruct | Qwen2ForCausalLM | ✅ |
-| Qwen/Qwen2.5-1.5B-Instruct | Qwen2ForCausalLM | 🟨 |
-
-✅ Runs and optimized.
-🟨 Runs and correct but not optimized to green yet.
-❌ Does not pass accuracy test or does not run.
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index bd14bbb9ab662..404519f887dc6 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -79,7 +79,9 @@ To make your model compatible with the Transformers modeling backend, it needs:
1. Add `is_causal = False` to `MyAttention`.
- If your model is mixture-of-experts (MoE):
1. Your sparse MoE block must have an attribute called `experts`.
- 2. The class of `experts` (`MyExperts`) must inherit from `nn.ModuleList`.
+ 2. The class of `experts` (`MyExperts`) must either:
+ - Inherit from `nn.ModuleList` (naive).
+ - Or contain all 3D `nn.Parameters` (packed).
3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.
2. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.
3. `MyModel` must contain `_supports_attention_backend = True`.
@@ -422,7 +424,7 @@ th {
| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ |
| `OLMoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
| `OLMo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
-| `OLMo3ForCausalLM` | OLMo3 | TBA | ✅︎ | ✅︎ |
+| `OLMo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |
| `OLMoEForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |
| `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ |
@@ -434,6 +436,7 @@ th {
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
| `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ |
| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | | ✅︎ |
+| `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | | ✅︎ |
| `QWenLMHeadModel` | Qwen | `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc. | ✅︎ | ✅︎ |
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ |
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ |
@@ -698,6 +701,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
| `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
| `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
+| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
| `Ovis` | Ovis2, Ovis1.6 | T + I+ | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
| `Ovis2_5` | Ovis2.5 | T + I+ + V | `AIDC-AI/Ovis2.5-9B`, etc. | | |
| `PaddleOCRVLForConditionalGeneration` | Paddle-OCR | T + I+ | `PaddlePaddle/PaddleOCR-VL`, etc. | | |
diff --git a/docs/serving/openai_compatible_server.md b/docs/serving/openai_compatible_server.md
index 821628e6e3174..23df3963823aa 100644
--- a/docs/serving/openai_compatible_server.md
+++ b/docs/serving/openai_compatible_server.md
@@ -293,7 +293,7 @@ and passing a list of `messages` in the request. Refer to the examples below for
base_url="http://localhost:8000/v1",
api_key="EMPTY",
)
- image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
response = create_chat_embeddings(
client,
diff --git a/docs/serving/parallelism_scaling.md b/docs/serving/parallelism_scaling.md
index 14cd3b057791c..a32840ea73b9a 100644
--- a/docs/serving/parallelism_scaling.md
+++ b/docs/serving/parallelism_scaling.md
@@ -118,14 +118,16 @@ The common practice is to set the tensor parallel size to the number of GPUs in
```bash
vllm serve /path/to/the/model/in/the/container \
--tensor-parallel-size 8 \
- --pipeline-parallel-size 2
+ --pipeline-parallel-size 2 \
+ --distributed-executor-backend ray
```
Alternatively, you can set `tensor_parallel_size` to the total number of GPUs in the cluster:
```bash
vllm serve /path/to/the/model/in/the/container \
- --tensor-parallel-size 16
+ --tensor-parallel-size 16 \
+ --distributed-executor-backend ray
```
## Optimizing network communication for tensor parallelism
diff --git a/docs/usage/reproducibility.md b/docs/usage/reproducibility.md
index d8a1943209c1e..a8e49d0a3398f 100644
--- a/docs/usage/reproducibility.md
+++ b/docs/usage/reproducibility.md
@@ -1,24 +1,23 @@
# Reproducibility
-vLLM does not guarantee the reproducibility of the results by default, for the sake of performance. You need to do the following to achieve
+vLLM does not guarantee the reproducibility of the results by default, for the sake of performance. To achieve
reproducible results:
-- For V1: Turn off multiprocessing to make the scheduling deterministic by setting `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
-- For V0: Set the global seed (see below).
+- In offline mode, you can either set `VLLM_ENABLE_V1_MULTIPROCESSING=0` which makes scheduling deterministic,
+ or enable [batch invariance](../features/batch_invariance.md) to make the outputs insensitive to scheduling.
+- In online mode, you can only enable [batch invariance](../features/batch_invariance.md).
Example: [examples/offline_inference/reproducibility.py](../../examples/offline_inference/reproducibility.py)
!!! warning
- Applying the above settings [changes the random state in user code](#locality-of-random-state).
+ Setting `VLLM_ENABLE_V1_MULTIPROCESSING=0` will change the random state of user code
+ (i.e. the code that constructs [LLM][vllm.LLM] class).
!!! note
Even with the above settings, vLLM only provides reproducibility
when it runs on the same hardware and the same vLLM version.
- Also, the online serving API (`vllm serve`) does not support reproducibility
- because it is almost impossible to make the scheduling deterministic in the
- online setting.
## Setting the global seed
@@ -26,27 +25,17 @@ The `seed` parameter in vLLM is used to control the random states for various ra
If a specific seed value is provided, the random states for `random`, `np.random`, and `torch.manual_seed` will be set accordingly.
-However, in some cases, setting the seed will also [change the random state in user code](#locality-of-random-state).
-
### Default Behavior
-In V0, the `seed` parameter defaults to `None`. When the `seed` parameter is `None`, the random states for `random`, `np.random`, and `torch.manual_seed` are not set. This means that each run of vLLM will produce different results if `temperature > 0`, as expected.
-
In V1, the `seed` parameter defaults to `0` which sets the random state for each worker, so the results will remain consistent for each vLLM run even if `temperature > 0`.
+It is impossible to un-specify a seed for V1 because different workers need to sample the same outputs
+for workflows such as speculative decoding. For more information, see:
+
!!! note
- It is impossible to un-specify a seed for V1 because different workers need to sample the same outputs
- for workflows such as speculative decoding.
-
- For more information, see:
+ The random state in user code (i.e. the code that constructs [LLM][vllm.LLM] class) is updated by vLLM
+ only if the workers are run in the same process as user code, i.e.: `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
-### Locality of random state
-
-The random state in user code (i.e. the code that constructs [LLM][vllm.LLM] class) is updated by vLLM under the following conditions:
-
-- For V0: The seed is specified.
-- For V1: The workers are run in the same process as user code, i.e.: `VLLM_ENABLE_V1_MULTIPROCESSING=0`.
-
-By default, these conditions are not active so you can use vLLM without having to worry about
-accidentally making deterministic subsequent operations that rely on random state.
+ By default, `VLLM_ENABLE_V1_MULTIPROCESSING=1` so you can use vLLM without having to worry about
+ accidentally making deterministic subsequent operations that rely on random state.
diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md
index 8d8a9e0f50805..5f647aafd61d4 100644
--- a/docs/usage/v1_guide.md
+++ b/docs/usage/v1_guide.md
@@ -2,11 +2,9 @@
!!! announcement
- We have started the process of deprecating V0. Please read [RFC #18571](https://github.com/vllm-project/vllm/issues/18571) for more details.
+ We have fully deprecated V0. Please read [RFC #18571](https://github.com/vllm-project/vllm/issues/18571) for more details.
-V1 is now enabled by default for all supported use cases, and we will gradually enable it for every use case we plan to support. Please share any feedback on [GitHub](https://github.com/vllm-project/vllm) or in the [vLLM Slack](https://inviter.co/vllm-slack).
-
-## Why vLLM V1?
+ If you have a use case that works on V0 Engine but not V1, please share it on [GitHub](https://github.com/vllm-project/vllm) or in the [vLLM Slack](https://inviter.co/vllm-slack).
vLLM V0 successfully supported a wide range of models and hardware, but as new features were developed independently, the system grew increasingly complex. This complexity made it harder to integrate new capabilities and introduced technical debt, revealing the need for a more streamlined and unified design.
@@ -32,16 +30,44 @@ Upgrade to vLLM’s Core Architecture](https://blog.vllm.ai/2025/01/27/v1-alpha-
This living user guide outlines a few known **important changes and limitations** introduced by vLLM V1. The team has been working actively to bring V1 as the default engine, therefore this guide will be updated constantly as more features get supported on vLLM V1.
-## Current Status
+## Differences from V0
-For each item, our progress towards V1 support falls into one of the following states:
+This section lists some differences in behavior between V0 and V1.
-- **🚀 Optimized**: Nearly fully optimized, with no further work currently planned.
-- **🟢 Functional**: Fully operational, with ongoing optimizations.
-- **🚧 WIP**: Under active development.
-- **🟡 Planned**: Scheduled for future implementation (some may have open PRs/RFCs).
-- **🟠 Delayed**: Temporarily dropped in V1 but planned to be re-introduced later.
-- **🔴 Deprecated**: Not planned for V1 unless there is strong demand.
+### Chunked Prefill
+
+Chunked prefill is enabled by default whenever possible, unlike in V0 where it was conditionally enabled based on model characteristics.
+
+### CUDA Graphs
+
+CUDA graph capture takes up more memory in V1 than in V0.
+
+### Semantic Changes to Logprobs
+
+#### Logprobs Calculation
+
+By default, logprobs in V1 are now returned immediately once computed from the model’s raw output (i.e.
+before applying any logits post-processing such as temperature scaling or penalty
+adjustments). As a result, the returned logprobs do not reflect the final adjusted
+probabilities used during sampling.
+
+You can adjust this behavior by setting the `--logprobs-mode` flag.
+Four modes are supported: `raw_logprobs` (default), `processed_logprobs`, `raw_logits`, `processed_logits`.
+Raw means the values before applying any logit processors, like bad words.
+Processed means the values after applying all processors, including temperature and top_k/top_p.
+
+#### Prompt Logprobs with Prefix Caching
+
+While V1 supports passing prompt logprobs with prefix caching enabled, it no longer caches the logprobs.
+For a request requiring prompt logprobs, the engine will ignore the prefix cache and recompute the prefill of full prompt to generate the logprobs.
+
+## Feature Support
+
+For each item, its support in vLLM V1 falls into one of the following states:
+
+- **🟢 Functional**: Fully operational with optimizations comparable to or better than V0.
+- **🟡 In Progress**: Planned to be in vLLM V1, with open PRs/RFCs.
+- **🔴 Removed**: Dropped from vLLM V1. Will only consider re-introducing if there is strong demand.
!!! note
vLLM V1’s unified scheduler treats both prompt and output tokens the same
@@ -57,13 +83,13 @@ based on assigned priority, with FCFS as a tie-breaker), configurable via the
### Hardware
-| Hardware | Status |
-|------------|-----------------------------------------------|
-| **NVIDIA** | 🚀 |
-| **AMD** | 🟢 |
+| Hardware | Status |
+|------------------|-----------------------------------------------|
+| **NVIDIA** | 🟢 |
+| **AMD** | 🟢 |
| **INTEL GPU** | 🟢 |
-| **TPU** | 🟢 |
-| **CPU** | 🟢 (x86\_64/aarch64) 🟡 (MacOS) |
+| **TPU** | 🟢 |
+| **CPU** | 🟢 |
!!! note
@@ -78,23 +104,21 @@ based on assigned priority, with FCFS as a tie-breaker), configurable via the
### Models
-| Model Type | Status |
-|-----------------------------|------------------------------------------------------------------------------------|
-| **Decoder-only Models** | 🚀 Optimized |
-| **Encoder-Decoder Models** | 🟢 Whisper only |
-| **Embedding Models** | 🟢 Functional |
-| **Mamba Models** | 🟢 (Mamba-2), 🟢 (Mamba-1) |
-| **Multimodal Models** | 🟢 Functional |
+| Model Type | Status |
+|-----------------------------|-------------------------------------------------------------------------|
+| **Decoder-only Models** | 🟢 |
+| **Encoder-Decoder Models** | 🟢 (Whisper), 🔴 (Others) |
+| **Pooling Models** | 🟢 |
+| **Mamba Models** | 🟢 |
+| **Multimodal Models** | 🟢 |
See below for the status of models that are not yet supported or have more features planned in V1.
-#### Embedding Models
+#### Pooling Models
-The initial basic support is now functional.
+Now fully supported, with prefix caching and chunked prefill newly available for last-pooling models.
-Later, we will consider using [hidden states processor](https://github.com/vllm-project/vllm/issues/12249),
-which is based on [global logits processor](https://github.com/vllm-project/vllm/pull/13360)
-to enable simultaneous generation and embedding using the same engine instance in V1.
+We are working on enabling prefix caching and chunked prefill for more categories of pooling models.
#### Mamba Models
@@ -112,24 +136,25 @@ Please note that prefix caching is not yet supported for any of the above models
Whisper is supported. Other models requiring cross-attention between separate
encoder and decoder (e.g., `BartForConditionalGeneration`,
-`MllamaForConditionalGeneration`) are not supported.
+`MllamaForConditionalGeneration`) are no longer supported.
### Features
| Feature | Status |
|---------------------------------------------|-----------------------------------------------------------------------------------|
-| **Prefix Caching** | 🚀 Optimized |
-| **Chunked Prefill** | 🚀 Optimized |
-| **LoRA** | 🚀 Optimized |
+| **Prefix Caching** | 🟢 Functional |
+| **Chunked Prefill** | 🟢 Functional |
+| **LoRA** | 🟢 Functional |
| **Logprobs Calculation** | 🟢 Functional |
-| **FP8 KV Cache** | 🟢 Functional on Hopper devices () |
-| **Spec Decode** | 🚀 Optimized |
-| **Prompt Logprobs with Prefix Caching** | 🟡 Planned ([RFC #13414](https://github.com/vllm-project/vllm/issues/13414)) |
+| **FP8 KV Cache** | 🟢 Functional |
+| **Spec Decode** | 🟢 Functional |
+| **Prompt Logprobs with Prefix Caching** | 🟢 Functional |
| **Structured Output Alternative Backends** | 🟢 Functional |
-| **Request-level Structured Output Backend** | 🔴 Deprecated |
-| **best_of** | 🔴 Deprecated ([RFC #13361](https://github.com/vllm-project/vllm/issues/13361)) |
-| **Per-Request Logits Processors** | 🔴 Deprecated ([RFC #13360](https://github.com/vllm-project/vllm/pull/13360)) |
-| **GPU <> CPU KV Cache Swapping** | 🔴 Deprecated |
+| **Concurrent Partial Prefills** | 🟡 [In Progress](https://github.com/vllm-project/vllm/issues/14003) |
+| **best_of** | 🔴 [Removed](https://github.com/vllm-project/vllm/issues/13361) |
+| **Per-Request Logits Processors** | 🔴 [Removed](https://github.com/vllm-project/vllm/pull/13360) |
+| **GPU <> CPU KV Cache Swapping** | 🔴 Removed |
+| **Request-level Structured Output Backend** | 🔴 Removed |
!!! note
@@ -139,38 +164,17 @@ encoder and decoder (e.g., `BartForConditionalGeneration`,
prefix caching, and speculative decoding without a strict separation between prefill
and decode phases.
-#### Semantic Changes to Logprobs
+#### Removed Features
-vLLM V1 supports logprobs and prompt logprobs. However, there are some important semantic
-differences compared to V0:
-
-##### Logprobs Calculation
-
-By default, logprobs in V1 are now returned immediately once computed from the model’s raw output (i.e.
-before applying any logits post-processing such as temperature scaling or penalty
-adjustments). As a result, the returned logprobs do not reflect the final adjusted
-probabilities used during sampling.
-
-You can adjust this behavior by setting the `--logprobs-mode` flag.
-Four modes are supported: `raw_logprobs` (default), `processed_logprobs`, `raw_logits`, `processed_logits`.
-Raw means the values before applying any logit processors, like bad words.
-Processed means the values after applying all processors, including temperature and top_k/top_p.
-
-##### Prompt Logprobs with Prefix Caching
-
-Logprobs are not cached. For a request requiring prompt logprobs, the engine will ignore the prefix cache and recompute the prefill of full prompt to generate the logprobs.
-
-#### Deprecated Features
-
-As part of the major architectural rework in vLLM V1, several legacy features have been deprecated.
+As part of the major architectural rework in vLLM V1, several legacy features have been removed.
##### Sampling features
-- **best_of**: This feature has been deprecated due to limited usage. See details at [RFC #13361](https://github.com/vllm-project/vllm/issues/13361).
+- **best_of**: This feature has been removed due to limited usage. See details at [RFC #13361](https://github.com/vllm-project/vllm/issues/13361).
- **Per-Request Logits Processors**: In V0, users could pass custom
processing functions to adjust logits on a per-request basis. In vLLM V1, this
- feature has been deprecated. Instead, the design is moving toward supporting **global logits
- processors**, a feature the team is actively working on for future releases. See details at [RFC #13360](https://github.com/vllm-project/vllm/pull/13360).
+ feature has been removed. Instead, we now support **global logits processors**
+ which are set at startup time, see [RFC #17799](https://github.com/vllm-project/vllm/issues/17799).
##### KV Cache features
@@ -179,4 +183,4 @@ to handle request preemptions.
##### Structured Output features
-- **Request-level Structured Output Backend**: Deprecated, alternative backends (outlines, guidance) with fallbacks is supported now.
+- **Request-level Structured Output Backend**: Removed; alternative backends (outlines, guidance) with fallbacks are supported now.
diff --git a/examples/offline_inference/context_extension.py b/examples/offline_inference/context_extension.py
index df39e4c25d5c8..67d33e1881ee9 100644
--- a/examples/offline_inference/context_extension.py
+++ b/examples/offline_inference/context_extension.py
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This script demonstrates how to extend the context length
-of a Qwen model using the YARN method (rope_scaling)
+of a Qwen model using the YARN method (rope_parameters)
and run a simple chat example.
Usage:
@@ -19,8 +19,8 @@ def create_llm():
# Use yarn to extend context
hf_overrides = {
- "rope_theta": rope_theta,
- "rope_scaling": {
+ "rope_parameters": {
+ "rope_theta": rope_theta,
"rope_type": "yarn",
"factor": factor,
"original_max_position_embeddings": original_max_position_embeddings,
diff --git a/examples/offline_inference/profiling_tpu/README.md b/examples/offline_inference/profiling_tpu/README.md
deleted file mode 100644
index 8c9c1c92b6764..0000000000000
--- a/examples/offline_inference/profiling_tpu/README.md
+++ /dev/null
@@ -1,70 +0,0 @@
-# vLLM TPU Profiling
-
-This script is used to profile the TPU performance of vLLM for specific prefill or decode token shapes.
-
-Note: an actual running server is a mix of both prefill of many shapes and decode of many shapes.
-
-We assume you are on a TPU already (this was tested on TPU v6e) and have installed vLLM according to the [Google TPU installation guide](https://docs.vllm.ai/en/latest/getting_started/installation/google_tpu.html).
-
-> In all examples below, we run several warmups before (so `--enforce-eager` is okay)
-
-## Profile Examples
-
-### Generate Prefill Trace
-
-This example runs Qwen/Qwen2.5-7B-Instruct with a single request of 1024 input tokens. This is set up in attempt to profile just the prefill time and operations.
-
-```bash
-export XLA_HLO_DEBUG=1
-export MODEL=Qwen/Qwen2.5-7B-Instruct
-export VLLM_TPU_PROFILE_DURATION_MS=3000
-export VLLM_TPU_PROFILE_DELAY_MS=0
-
-python3 profiling.py \
- --model $MODEL \
- --input-len 1024 --output-len 1 \
- --batch-size 1 --enforce-eager \
- --max-model-len 2048 \
- --tensor-parallel-size 1 \
- --profile-result-dir profiles
-```
-
-### Generate Decode Trace
-
-This example runs Llama 3.1 70B with a batch of 32 requests where each has 1 input token and 128 output tokens. This is set up in attempt to profile just the 32 decodes running in parallel by having an extremely small prefill of 1 token and setting `VLLM_TPU_PROFILE_DELAY_MS=1000` to skip the first second of inference (hopefully prefill).
-
-```bash
-export XLA_HLO_DEBUG=1
-export MODEL=meta-llama/Llama-3.1-70B-Instruct
-export VLLM_TPU_PROFILE_DURATION_MS=2000
-export VLLM_TPU_PROFILE_DELAY_MS=1000
-
-rm -rf ~/.cache/vllm/xla_cache
-python3 profiling.py \
- --model $MODEL \
- --input-len 1 \
- --output-len 128 \
- --batch-size 32 \
- --enforce-eager \
- --profile-result-dir profiles \
- --max-model-len 2048 --tensor-parallel-size 8
-```
-
-## Visualizing the profiles
-
-Once you have collected your profiles with this script, you can visualize them using [TensorBoard](https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm).
-
-Here are most likely the dependencies you need to install:
-
-```bash
-pip install tensorflow-cpu \
- tensorboard-plugin-profile \
- etils \
- importlib_resources
-```
-
-Then you just need to point TensorBoard to the directory where you saved the profiles and visit `http://localhost:6006/` in your browser:
-
-```bash
-tensorboard --logdir profiles/ --port 6006
-```
diff --git a/examples/offline_inference/profiling_tpu/profiling.py b/examples/offline_inference/profiling_tpu/profiling.py
deleted file mode 100644
index 3b127e4fd29df..0000000000000
--- a/examples/offline_inference/profiling_tpu/profiling.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import dataclasses
-import os
-import time
-
-import numpy as np
-import torch_xla.debug.profiler as xp
-from tqdm import tqdm
-
-from vllm import LLM, SamplingParams
-from vllm.engine.arg_utils import EngineArgs
-from vllm.inputs import PromptType
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-DURATION_MS = int(os.getenv("VLLM_TPU_PROFILE_DURATION_MS", 3000))
-DELAY_MS = int(os.getenv("VLLM_TPU_PROFILE_DELAY_MS", 0))
-
-
-def main(args: argparse.Namespace):
- print(args)
-
- engine_args = EngineArgs.from_cli_args(args)
- llm = LLM(**dataclasses.asdict(engine_args))
- server = xp.start_server(9012) # noqa: F841
-
- sampling_params = SamplingParams(
- temperature=0.0,
- ignore_eos=True,
- max_tokens=args.output_len,
- )
- print(sampling_params)
- dummy_prompt_token_ids = np.random.randint(
- 10000, size=(args.batch_size, args.input_len)
- )
- dummy_prompts: list[PromptType] = [
- {"prompt_token_ids": batch} for batch in dummy_prompt_token_ids.tolist()
- ]
-
- def run_to_completion():
- start_time = time.perf_counter()
- llm.generate(dummy_prompts, sampling_params=sampling_params, use_tqdm=False)
- end_time = time.perf_counter()
- latency = end_time - start_time
- return latency
-
- # Warmup
- print("Warming up...")
- warmup_latencies = []
- for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
- warmup_latencies.append(run_to_completion())
- print(f"Average warmup latency: {np.mean(warmup_latencies):.4f}s")
-
- # Profile
- profile_dir = args.profile_result_dir
- print(f"Profiling (results will be saved to '{profile_dir}')...")
- # Enable tracing on server
- xp.trace_detached(
- "localhost:9012", profile_dir, delay_ms=DELAY_MS, duration_ms=DURATION_MS
- )
- if DELAY_MS == 0:
- time.sleep(1.0)
- profile_latencies = []
- for _ in tqdm(range(args.num_iters), desc="Profile iterations"):
- profile_latencies.append(run_to_completion())
- print(f"Average profile latency: {np.mean(profile_latencies):.4f}s")
-
- return
-
-
-def parse_args():
- parser = FlexibleArgumentParser(
- description="Benchmark the latency of processing a single batch of "
- "requests till completion."
- )
- parser.add_argument("--input-len", type=int, default=32)
- parser.add_argument("--output-len", type=int, default=128)
- parser.add_argument("--batch-size", type=int, default=8)
- parser.add_argument(
- "--num-iters-warmup",
- type=int,
- default=5,
- help="Number of iterations to run for warmup.",
- )
- parser.add_argument(
- "--num-iters",
- type=int,
- default=1,
- help="Number of iterations to run for profiling.",
- )
- parser.add_argument(
- "--profile-result-dir",
- type=str,
- default="profiles",
- help=(
- "path to save the pytorch profiler output. Can be visualized "
- "with ui.perfetto.dev or Tensorboard "
- "(https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm)."
- ),
- )
-
- parser = EngineArgs.add_cli_args(parser)
- return parser.parse_args()
-
-
-if __name__ == "__main__":
- args = parse_args()
- main(args)
diff --git a/examples/offline_inference/reproducibility.py b/examples/offline_inference/reproducibility.py
index d909438b41042..72c1e841dca45 100644
--- a/examples/offline_inference/reproducibility.py
+++ b/examples/offline_inference/reproducibility.py
@@ -11,12 +11,11 @@ import random
from vllm import LLM, SamplingParams
-# V1 only: Turn off multiprocessing to make the scheduling deterministic.
+# Either:
+## Turn off multiprocessing to make the scheduling deterministic, or
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-
-# V0 only: Set the global seed. The default seed is None, which is
-# not reproducible.
-SEED = 42
+## Enable batch invariance to get consistent results regardless of scheduling.
+os.environ["VLLM_BATCH_INVARIANT"] = "1"
prompts = [
"Hello, my name is",
@@ -28,7 +27,7 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
def main():
- llm = LLM(model="facebook/opt-125m", seed=SEED)
+ llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params)
print("-" * 50)
for output in outputs:
diff --git a/examples/offline_inference/rlhf.py b/examples/offline_inference/rlhf.py
index 0c09e603271de..6f05968ce065e 100644
--- a/examples/offline_inference/rlhf.py
+++ b/examples/offline_inference/rlhf.py
@@ -62,7 +62,7 @@ ray.init()
# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
# Learn more about Ray placement groups:
-# https://docs.ray.io/en/latest/placement-groups.html
+# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
ray.get(pg_inference.ready())
scheduling_inference = PlacementGroupSchedulingStrategy(
diff --git a/examples/offline_inference/rlhf_online_quant.py b/examples/offline_inference/rlhf_online_quant.py
new file mode 100644
index 0000000000000..2d98ad22c589e
--- /dev/null
+++ b/examples/offline_inference/rlhf_online_quant.py
@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.
+
+The script separates training and inference workloads onto distinct GPUs
+so that Ray can manage process placement and inter-process communication.
+A Hugging Face Transformer model occupies GPU 0 for training, whereas a
+tensor-parallel vLLM inference engine occupies GPU 1–2.
+
+The example performs the following steps:
+
+* Load the training model on GPU 0.
+* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
+ and Ray placement groups.
+* Generate text from a list of prompts using the inference engine.
+* Update the weights of the training model and broadcast the updated weights
+ to the inference engine by using a Ray collective RPC group. Note that
+ for demonstration purposes we simply zero out the weights.
+
+For a production-ready implementation that supports multiple training and
+inference replicas, see the OpenRLHF framework:
+https://github.com/OpenRLHF/OpenRLHF
+
+This example assumes a single-node cluster with three GPUs, but Ray
+supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
+workloads. Residual GPU activity interferes with vLLM memory profiling and
+causes unexpected behavior.
+"""
+
+import json
+import os
+
+import ray
+import torch
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from rlhf_utils import stateless_init_process_group
+from torchao.core.config import config_to_dict
+from torchao.quantization import (
+ Float8DynamicActivationFloat8WeightConfig,
+ PerRow,
+)
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.utils.network_utils import get_ip, get_open_port
+
+
+class MyLLM(LLM):
+ """Configure the vLLM worker for Ray placement group execution."""
+
+ def __init__(self, *args, **kwargs):
+ # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
+ # so that vLLM can manage its own device placement within the worker.
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ super().__init__(*args, **kwargs)
+
+
+# Load the OPT-125M model onto GPU 0 for the training workload.
+train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
+train_model.to("cuda:0")
+
+# Initialize Ray and set the visible devices. The vLLM engine will
+# be placed on GPUs 1 and 2.
+os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
+ray.init()
+
+# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
+# Learn more about Ray placement groups:
+# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
+pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
+ray.get(pg_inference.ready())
+scheduling_inference = PlacementGroupSchedulingStrategy(
+ placement_group=pg_inference,
+ placement_group_capture_child_tasks=True,
+ placement_group_bundle_index=0,
+)
+
+# Launch the vLLM inference engine. The `enforce_eager` flag reduces
+# start-up latency.
+
+# generate torchao quantization config for RL rollout
+# see https://github.com/vllm-project/vllm/pull/23014 for instructions to
+# use serialized config files instead of passing around json string
+config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
+
+json_str = json.dumps(config_to_dict(config))
+
+llm = ray.remote(
+ num_cpus=0,
+ num_gpus=0,
+ scheduling_strategy=scheduling_inference,
+)(MyLLM).remote(
+ model="facebook/opt-125m",
+ hf_overrides={"quantization_config_dict_json": json_str},
+ enforce_eager=True,
+ worker_extension_cls="rlhf_utils.WorkerExtension",
+ tensor_parallel_size=2,
+ distributed_executor_backend="ray",
+)
+
+# Generate text from the prompts.
+prompts = [
+ "Hello, my name is",
+ "The president of the United States is",
+ "The capital of France is",
+ "The future of AI is",
+]
+
+sampling_params = SamplingParams(temperature=0)
+
+outputs = ray.get(llm.generate.remote(prompts, sampling_params))
+
+print("-" * 50)
+for output in outputs:
+ prompt = output.prompt
+ generated_text = output.outputs[0].text
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+ print("-" * 50)
+
+# Set up the communication channel between the training process and the
+# inference engine.
+master_address = get_ip()
+master_port = get_open_port()
+
+handle = llm.collective_rpc.remote(
+ "init_weight_update_group", args=(master_address, master_port, 1, 3)
+)
+
+model_update_group = stateless_init_process_group(
+ master_address, master_port, 0, 3, torch.device("cuda:0")
+)
+ray.get(handle)
+
+# Simulate a training step by zeroing out all model weights.
+# In a real RLHF training loop the weights would be updated using the gradient
+# from an RL objective such as PPO on a reward model.
+for name, p in train_model.named_parameters():
+ p.data.zero_()
+
+# Synchronize the updated weights to the inference engine.
+for name, p in train_model.named_parameters():
+ dtype_name = str(p.dtype).split(".")[-1]
+ handle = llm.collective_rpc.remote(
+ "update_weight", args=(name, dtype_name, p.shape)
+ )
+ model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
+ ray.get(handle)
+
+# Verify that the inference weights have been updated.
+assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
+
+# Generate text with the updated model. The output is expected to be nonsense
+# because the weights are zero.
+outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
+print("-" * 50)
+for output in outputs_updated:
+ prompt = output.prompt
+ generated_text = output.outputs[0].text
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+ print("-" * 50)
diff --git a/examples/offline_inference/rlhf_utils.py b/examples/offline_inference/rlhf_utils.py
index 13def88439ef2..5c0787b8778d6 100644
--- a/examples/offline_inference/rlhf_utils.py
+++ b/examples/offline_inference/rlhf_utils.py
@@ -30,8 +30,8 @@ class WorkerExtension:
"""
The class for vLLM's worker to inherit from.
By defining an extension class, the code can work no matter what is
- the underlying worker class. This way, the code can be compatible
- with both vLLM V0 and V1.
+ the underlying worker class.
+
NOTE: we define this class in a separate module, and the main module
should pass the full qualified name as `worker_extension_cls` argument.
"""
@@ -96,8 +96,8 @@ class ColocateWorkerExtension:
"""
The class for vLLM's worker to inherit from, in the colocate setting.
By defining an extension class, the code can work no matter what is
- the underlying worker class. This way, the code can be compatible
- with both vLLM V0 and V1.
+ the underlying worker class.
+
NOTE: we define this class in a separate module, and the main module
should pass the full qualified name as `worker_extension_cls` argument.
"""
diff --git a/examples/offline_inference/save_sharded_state.py b/examples/offline_inference/save_sharded_state.py
index e25f46b126e6f..88ee48b98bff6 100644
--- a/examples/offline_inference/save_sharded_state.py
+++ b/examples/offline_inference/save_sharded_state.py
@@ -67,22 +67,9 @@ def main(args):
Path(args.output).mkdir(exist_ok=True)
# Dump worker states to output directory
- # Check which engine version is being used
- is_v1_engine = hasattr(llm.llm_engine, "engine_core")
-
- if is_v1_engine:
- # For V1 engine, we need to use engine_core.save_sharded_state
- print("Using V1 engine save path")
- llm.llm_engine.engine_core.save_sharded_state(
- path=args.output, pattern=args.file_pattern, max_size=args.max_file_size
- )
- else:
- # For V0 engine
- print("Using V0 engine save path")
- model_executor = llm.llm_engine.model_executor
- model_executor.save_sharded_state(
- path=args.output, pattern=args.file_pattern, max_size=args.max_file_size
- )
+ llm.llm_engine.engine_core.save_sharded_state(
+ path=args.output, pattern=args.file_pattern, max_size=args.max_file_size
+ )
# Copy metadata files to output directory
for file in os.listdir(model_path):
diff --git a/examples/offline_inference/spec_decode.py b/examples/offline_inference/spec_decode.py
index 3cdc3b245b72a..67a0732459709 100644
--- a/examples/offline_inference/spec_decode.py
+++ b/examples/offline_inference/spec_decode.py
@@ -158,11 +158,7 @@ def main(args):
print(f"generated text: {output.outputs[0].text}")
print("-" * 50)
- try:
- metrics = llm.get_metrics()
- except AssertionError:
- print("Metrics are not supported in the V0 engine.")
- return
+ metrics = llm.get_metrics()
total_num_output_tokens = sum(
len(output.outputs[0].token_ids) for output in outputs
diff --git a/examples/offline_inference/tpu.py b/examples/offline_inference/tpu.py
deleted file mode 100644
index 0093b63b0b1f3..0000000000000
--- a/examples/offline_inference/tpu.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import os
-
-from vllm import LLM, SamplingParams
-
-prompts = [
- "A robot may not injure a human being",
- "It is only with the heart that one can see rightly;",
- "The greatest glory in living lies not in never falling,",
-]
-answers = [
- " or, through inaction, allow a human being to come to harm.",
- " what is essential is invisible to the eye.",
- " but in rising every time we fall.",
-]
-N = 1
-# Currently, top-p sampling is disabled. `top_p` should be 1.0.
-sampling_params = SamplingParams(temperature=0, top_p=1.0, n=N, max_tokens=16)
-
-
-def main():
- parser = argparse.ArgumentParser(description="TPU offline inference example")
- parser.add_argument("--use-spmd", action="store_true", help="Enable SPMD mode")
- args = parser.parse_args()
-
- llm_args = {
- "model": "Qwen/Qwen2-1.5B-Instruct",
- "max_num_batched_tokens": 64,
- "max_num_seqs": 4,
- "max_model_len": 128,
- }
- if args.use_spmd:
- os.environ["VLLM_XLA_USE_SPMD"] = "1"
- # Can only hardcode the number of chips for now.
- # calling xr.global_runtime_device_count() beforeing init SPMD env in
- # torch_xla will mess up the distributed env.
- llm_args["tensor_parallel_size"] = 8
- # Use Llama, for num_kv_heads = 8.
- llm_args["model"] = "meta-llama/Llama-3.1-8B-Instruct"
-
- # Set `enforce_eager=True` to avoid ahead-of-time compilation.
- # In real workloads, `enforce_eager` should be `False`.
- llm = LLM(**llm_args)
- outputs = llm.generate(prompts, sampling_params)
- print("-" * 50)
- for output, answer in zip(outputs, answers):
- prompt = output.prompt
- generated_text = output.outputs[0].text
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
- assert generated_text.startswith(answer)
- print("-" * 50)
-
-
-if __name__ == "__main__":
- main()
diff --git a/examples/offline_inference/vision_language_multi_image.py b/examples/offline_inference/vision_language_multi_image.py
index d6e169548f15b..301265d4e17f7 100644
--- a/examples/offline_inference/vision_language_multi_image.py
+++ b/examples/offline_inference/vision_language_multi_image.py
@@ -1110,6 +1110,7 @@ def load_r_vl(question: str, image_urls: list[str]) -> ModelRequestData:
model=model_name,
max_model_len=16384,
max_num_seqs=16,
+ trust_remote_code=True,
limit_mm_per_prompt={"image": len(image_urls)},
)
diff --git a/examples/offline_inference/vision_language_pooling.py b/examples/offline_inference/vision_language_pooling.py
index 63d85d5d9eef5..530aad4bc031c 100644
--- a/examples/offline_inference/vision_language_pooling.py
+++ b/examples/offline_inference/vision_language_pooling.py
@@ -266,7 +266,7 @@ def get_query(modality: QueryModality):
return ImageQuery(
modality="image",
image=fetch_image(
- "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/American_Eskimo_Dog.jpg/360px-American_Eskimo_Dog.jpg" # noqa: E501
+ "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
),
)
@@ -275,7 +275,7 @@ def get_query(modality: QueryModality):
modality="text+image",
text="A cat standing in the snow.",
image=fetch_image(
- "https://upload.wikimedia.org/wikipedia/commons/thumb/b/b6/Felis_catus-cat_on_snow.jpg/179px-Felis_catus-cat_on_snow.jpg" # noqa: E501
+ "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg" # noqa: E501
),
)
diff --git a/examples/online_serving/disaggregated_prefill.sh b/examples/online_serving/disaggregated_prefill.sh
index d434e22b1ae88..cd2f2e44a4d69 100644
--- a/examples/online_serving/disaggregated_prefill.sh
+++ b/examples/online_serving/disaggregated_prefill.sh
@@ -24,7 +24,14 @@ cleanup() {
exit 0
}
-export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
+
+if [[ -z "${VLLM_HOST_IP:-}" ]]; then
+ export VLLM_HOST_IP=127.0.0.1
+ echo "Using default VLLM_HOST_IP=127.0.0.1 (override by exporting VLLM_HOST_IP before running this script)"
+else
+ echo "Using provided VLLM_HOST_IP=${VLLM_HOST_IP}"
+fi
+
# install quart first -- required for disagg prefill proxy serve
if python3 -c "import quart" &> /dev/null; then
@@ -38,7 +45,7 @@ fi
wait_for_server() {
local port=$1
timeout 1200 bash -c "
- until curl -s localhost:${port}/v1/completions > /dev/null; do
+ until curl -i localhost:${port}/v1/models > /dev/null; do
sleep 1
done" && return 0 || return 1
}
@@ -48,21 +55,23 @@ wait_for_server() {
# prefilling instance, which is the KV producer
CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL_NAME \
+ --host 0.0.0.0 \
--port 8100 \
--max-model-len 100 \
--gpu-memory-utilization 0.8 \
--trust-remote-code \
--kv-transfer-config \
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}' &
+ '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
-# decoding instance, which is the KV consumer
+# decoding instance, which is the KV consumer
CUDA_VISIBLE_DEVICES=1 vllm serve $MODEL_NAME \
+ --host 0.0.0.0 \
--port 8200 \
--max-model-len 100 \
--gpu-memory-utilization 0.8 \
--trust-remote-code \
--kv-transfer-config \
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}' &
+ '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":"1e10","kv_port":"14580","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8200","send_type":"PUT_ASYNC"}}' &
# wait until prefill and decode instances are ready
wait_for_server 8100
diff --git a/examples/online_serving/gradio_openai_chatbot_webserver.py b/examples/online_serving/gradio_openai_chatbot_webserver.py
index d5d0a07a29183..c76c60cc4472d 100644
--- a/examples/online_serving/gradio_openai_chatbot_webserver.py
+++ b/examples/online_serving/gradio_openai_chatbot_webserver.py
@@ -25,25 +25,17 @@ import gradio as gr
from openai import OpenAI
-def format_history_to_openai(history):
- history_openai_format = [
- {"role": "system", "content": "You are a great AI assistant."}
- ]
- for human, assistant in history:
- history_openai_format.append({"role": "user", "content": human})
- history_openai_format.append({"role": "assistant", "content": assistant})
- return history_openai_format
-
-
def predict(message, history, client, model_name, temp, stop_token_ids):
- # Format history to OpenAI chat format
- history_openai_format = format_history_to_openai(history)
- history_openai_format.append({"role": "user", "content": message})
+ messages = [
+ {"role": "system", "content": "You are a great AI assistant."},
+ *history,
+ {"role": "user", "content": message},
+ ]
# Send request to OpenAI API (vLLM server)
stream = client.chat.completions.create(
model=model_name,
- messages=history_openai_format,
+ messages=messages,
temperature=temp,
stream=True,
extra_body={
diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal.py b/examples/online_serving/openai_chat_completion_client_for_multimodal.py
index 520cbca003aa5..3d1259276998d 100644
--- a/examples/online_serving/openai_chat_completion_client_for_multimodal.py
+++ b/examples/online_serving/openai_chat_completion_client_for_multimodal.py
@@ -66,7 +66,7 @@ def run_text_only(model: str, max_completion_tokens: int) -> None:
# Single-image input inference
def run_single_image(model: str, max_completion_tokens: int) -> None:
## Use image url in the payload
- image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
chat_completion_from_url = client.chat.completions.create(
messages=[
{
diff --git a/examples/online_serving/openai_embedding_long_text/service.sh b/examples/online_serving/openai_embedding_long_text/service.sh
index 1577de85f7ff2..b5c92749466b0 100644
--- a/examples/online_serving/openai_embedding_long_text/service.sh
+++ b/examples/online_serving/openai_embedding_long_text/service.sh
@@ -22,7 +22,6 @@ API_KEY=${API_KEY:-"your-api-key"}
POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST
export VLLM_ENABLE_CHUNKED_PROCESSING=true
export CUDA_VISIBLE_DEVICES=2,3,4,5
-# export VLLM_ATTENTION_BACKEND=XFORMERS
echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing"
echo "=================================================================="
diff --git a/examples/online_serving/openai_responses_client.py b/examples/online_serving/openai_responses_client.py
new file mode 100644
index 0000000000000..b4eb24671507a
--- /dev/null
+++ b/examples/online_serving/openai_responses_client.py
@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Set up this example by starting a vLLM OpenAI-compatible server.
+Reasoning models can be used through the Responses API as seen here
+https://platform.openai.com/docs/api-reference/responses
+For example:
+vllm serve Qwen/Qwen3-8B --reasoning-parser qwen3
+
+"""
+
+from openai import OpenAI
+
+input_messages = [{"role": "user", "content": "What model are you?"}]
+
+
+def main():
+ base_url = "http://localhost:8000/v1"
+ client = OpenAI(base_url=base_url, api_key="empty")
+ model = "Qwen/Qwen3-8B" # get_first_model(client)
+ response = client.responses.create(
+ model=model,
+ input=input_messages,
+ )
+
+ for message in response.output:
+ if message.type == "reasoning":
+ # append reasoning message
+ input_messages.append(message)
+
+ response_2 = client.responses.create(
+ model=model,
+ input=input_messages,
+ )
+ print(response_2.output_text)
+ # I am Qwen, a large language model developed by Alibaba Cloud.
+ # I am designed to assist with a wide range of tasks, including
+ # answering questions, creating content, coding, and engaging in
+ # conversations. I can help with various topics and provide
+ # information or support in multiple languages. How can I assist you today?
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py b/examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py
index 261b810ce5d03..47c2c5030078c 100644
--- a/examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py
+++ b/examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py
@@ -21,7 +21,7 @@ from PIL import Image
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"
-image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
def create_chat_embeddings(
diff --git a/examples/online_serving/prometheus_grafana/README.md b/examples/online_serving/prometheus_grafana/README.md
index 5cd4dab5a8fa7..9615210a2ad80 100644
--- a/examples/online_serving/prometheus_grafana/README.md
+++ b/examples/online_serving/prometheus_grafana/README.md
@@ -46,7 +46,7 @@ Navigate to [`http://localhost:3000`](http://localhost:3000). Log in with the de
Navigate to [`http://localhost:3000/connections/datasources/new`](http://localhost:3000/connections/datasources/new) and select Prometheus.
-On Prometheus configuration page, we need to add the `Prometheus Server URL` in `Connection`. For this setup, Grafana and Prometheus are running in separate containers, but Docker creates DNS name for each containers. You can just use `http://prometheus:9090`.
+On Prometheus configuration page, we need to add the `Prometheus Server URL` in `Connection`. For this setup, Grafana and Prometheus are running in separate containers, but Docker creates DNS name for each container. You can just use `http://prometheus:9090`.
Click `Save & Test`. You should get a green check saying "Successfully queried the Prometheus API.".
diff --git a/format.sh b/format.sh
deleted file mode 100755
index 6ba93e0a19ba8..0000000000000
--- a/format.sh
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/bin/bash
-
-echo "vLLM linting system has been moved from format.sh to pre-commit hooks."
-echo "Please run 'pip install -r requirements/lint.txt', followed by"
-echo "'pre-commit install' to install the pre-commit hooks."
-echo "Then linters will run automatically before each commit."
\ No newline at end of file
diff --git a/requirements/common.txt b/requirements/common.txt
index ad92ba3ad8278..3f8cd588422d0 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -19,18 +19,18 @@ pillow # Required for image processing
prometheus-fastapi-instrumentator >= 7.0.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
-llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x"
+llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x" or platform_machine == "ppc64le"
outlines_core == 0.2.11
# required for outlines backend disk cache
diskcache == 5.6.3
lark == 1.2.2
-xgrammar == 0.1.25; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x"
+xgrammar == 0.1.27; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
typing_extensions >= 4.10
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
partial-json-parser # used for parsing partial JSON outputs
pyzmq >= 25.0.0
msgspec
-gguf >= 0.13.0
+gguf >= 0.17.0
mistral_common[image] >= 1.8.5
opencv-python-headless >= 4.11.0 # required for video IO
pyyaml
diff --git a/requirements/cpu-build.txt b/requirements/cpu-build.txt
index 331d02be6621e..81d429a5e5f8d 100644
--- a/requirements/cpu-build.txt
+++ b/requirements/cpu-build.txt
@@ -4,8 +4,9 @@ packaging>=24.2
setuptools>=77.0.3,<81.0.0
setuptools-scm>=8
--extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.8.0+cpu; platform_machine == "x86_64"
-torch==2.8.0; platform_machine == "ppc64le" or platform_machine == "aarch64" or platform_system == "Darwin"
+torch==2.8.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
+torch==2.9.0; platform_system == "Darwin"
+torch==2.8.0; platform_machine == "ppc64le" or platform_machine == "aarch64"
scons; platform_machine == "aarch64" # needed to build Arm Compute Library (ACL)
wheel
jinja2>=3.1.6
diff --git a/requirements/cpu.txt b/requirements/cpu.txt
index d11787df4d92b..e23d3286f3f78 100644
--- a/requirements/cpu.txt
+++ b/requirements/cpu.txt
@@ -22,7 +22,6 @@ datasets # for benchmark scripts
# Intel Extension for PyTorch, only for x86_64 CPUs
intel-openmp==2024.2.1; platform_machine == "x86_64"
-intel_extension_for_pytorch==2.8.0; platform_machine == "x86_64"
triton==3.2.0; platform_machine == "x86_64" # Triton is required for torch 2.6+cpu, as it is imported in torch.compile.
# Use this to gather CPU info and optimize based on ARM Neoverse cores
diff --git a/requirements/cuda.txt b/requirements/cuda.txt
index d63fe9e1e77c1..15e8aadc56f47 100644
--- a/requirements/cuda.txt
+++ b/requirements/cuda.txt
@@ -9,6 +9,5 @@ torch==2.9.0
torchaudio==2.9.0
# These must be updated alongside torch
torchvision==0.24.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
-xformers==0.0.33.post1; platform_system == 'Linux' and platform_machine == 'x86_64' # Requires PyTorch >= 2.9
# FlashInfer should be updated together with the Dockerfile
flashinfer-python==0.5.2
diff --git a/requirements/rocm-test.txt b/requirements/rocm-test.txt
index 432e11977872d..2d57e7e167869 100644
--- a/requirements/rocm-test.txt
+++ b/requirements/rocm-test.txt
@@ -39,3 +39,9 @@ mteb[bm25s]>=1.38.11, <2
# Required for eval tests
lm-eval[api] @ git+https://github.com/EleutherAI/lm-evaluation-harness.git@206b7722158f58c35b7ffcd53b035fdbdda5126d
+
+# Required for multiprocessed tests that use spawn method
+multiprocess==0.70.16
+
+# Plugins test
+terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
diff --git a/requirements/rocm.txt b/requirements/rocm.txt
index 6f1cca90e5e2b..abbd33d6e1240 100644
--- a/requirements/rocm.txt
+++ b/requirements/rocm.txt
@@ -15,3 +15,4 @@ setuptools-scm>=8
runai-model-streamer[s3,gcs]==0.15.0
conch-triton-kernels==1.2.1
timm>=1.0.17
+fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@d6f998a03432b2452f8de2bb5cefb5af9795d459
diff --git a/requirements/test.in b/requirements/test.in
index 30d97e9b9c7d0..05f6bcca5c2c4 100644
--- a/requirements/test.in
+++ b/requirements/test.in
@@ -36,7 +36,7 @@ opencv-python-headless >= 4.11.0 # required for video test
datamodel_code_generator # required for minicpm3 test
# TODO: Use lm-eval[api]==0.4.10 once released
lm-eval[api] @ git+https://github.com/EleutherAI/lm-evaluation-harness.git@206b7722158f58c35b7ffcd53b035fdbdda5126d # required for model evaluation test
-mteb[bm25s]>=1.38.11, <2 # required for mteb test
+mteb[bm25s]>=2, <3 # required for mteb test
transformers==4.57.1
tokenizers==0.22.0
schemathesis>=3.39.15 # Required for openai schema test.
diff --git a/requirements/test.txt b/requirements/test.txt
index 3263b74c08797..bcd511660f85e 100644
--- a/requirements/test.txt
+++ b/requirements/test.txt
@@ -201,8 +201,6 @@ email-validator==2.2.0
# via pydantic
encodec==0.1.1
# via vocos
-eval-type-backport==0.2.2
- # via mteb
evaluate==0.4.3
# via lm-eval
fastapi==0.116.1
@@ -490,7 +488,7 @@ msgpack==1.1.0
# via
# librosa
# ray
-mteb==1.38.11
+mteb==2.1.2
# via -r requirements/test.in
multidict==6.1.0
# via
diff --git a/setup.py b/setup.py
index e9b36e2a2e037..8871b04d8fc46 100644
--- a/setup.py
+++ b/setup.py
@@ -74,18 +74,6 @@ def is_ninja_available() -> bool:
return which("ninja") is not None
-def is_url_available(url: str) -> bool:
- from urllib.request import urlopen
-
- status = None
- try:
- with urlopen(url) as f:
- status = f.status
- except Exception:
- return False
- return status == 200
-
-
class CMakeExtension(Extension):
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
super().__init__(name, sources=[], py_limited_api=True, **kwa)
@@ -299,6 +287,20 @@ class cmake_build_ext(build_ext):
os.makedirs(os.path.dirname(dst_file), exist_ok=True)
self.copy_file(file, dst_file)
+ if _is_cuda() or _is_hip():
+ # copy vllm/third_party/triton_kernels/**/*.py from self.build_lib
+ # to current directory so that they can be included in the editable
+ # build
+ print(
+ f"Copying {self.build_lib}/vllm/third_party/triton_kernels "
+ "to vllm/third_party/triton_kernels"
+ )
+ shutil.copytree(
+ f"{self.build_lib}/vllm/third_party/triton_kernels",
+ "vllm/third_party/triton_kernels",
+ dirs_exist_ok=True,
+ )
+
class precompiled_build_ext(build_ext):
"""Disables extension building when using precompiled binaries."""
@@ -519,28 +521,6 @@ def get_nvcc_cuda_version() -> Version:
return nvcc_cuda_version
-def get_gaudi_sw_version():
- """
- Returns the driver version.
- """
- # Enable console printing for `hl-smi` check
- output = subprocess.run(
- "hl-smi",
- shell=True,
- text=True,
- capture_output=True,
- env={"ENABLE_CONSOLE": "true"},
- )
- if output.returncode == 0 and output.stdout:
- return (
- output.stdout.split("\n")[2]
- .replace(" ", "")
- .split(":")[1][:-1]
- .split("-")[0]
- )
- return "0.0.0" # when hl-smi is not available
-
-
def get_vllm_version() -> str:
# Allow overriding the version. This is useful to build platform-specific
# wheels (e.g. CPU, TPU) without modifying the source.
@@ -633,6 +613,9 @@ ext_modules = []
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._moe_C"))
ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
+ # Optional since this doesn't get built (produce an .so file). This is just
+ # copying the relevant .py files from the source repository.
+ ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py
index 0cf1e85d4e8ee..521d6c33dd390 100644
--- a/tests/basic_correctness/test_basic_correctness.py
+++ b/tests/basic_correctness/test_basic_correctness.py
@@ -74,9 +74,6 @@ def test_models(
model_executor: str,
enable_prompt_embeds: bool,
) -> None:
- if backend == "XFORMERS" and model == "google/gemma-2-2b-it":
- pytest.skip(f"{backend} does not support gemma2 with full context length.")
-
with monkeypatch.context() as m:
m.setenv("VLLM_ATTENTION_BACKEND", backend)
diff --git a/tests/compile/README.md b/tests/compile/README.md
new file mode 100644
index 0000000000000..300a956860005
--- /dev/null
+++ b/tests/compile/README.md
@@ -0,0 +1,5 @@
+# compile test folder structure
+
+- `compile/test_*.py` : various unit tests meant for testing particular code path/features. Future tests are most likely added here. New test files added here will be included in CI automatically
+- `compile/fullgraph/` : full model tests, including all tests previously in compile/piecewise. These tests do not target particular features. New test files added here will be included in CI automatically
+- `compile/distributed/` : tests that require multiple GPUs. New test files added here will **NOT** be included in CI automatically as these tests generally need to be manually configured to run in runners with particular number/type of GPUs.
diff --git a/tests/compile/piecewise/__init__.py b/tests/compile/distributed/__init__.py
similarity index 100%
rename from tests/compile/piecewise/__init__.py
rename to tests/compile/distributed/__init__.py
diff --git a/tests/compile/test_async_tp.py b/tests/compile/distributed/test_async_tp.py
similarity index 99%
rename from tests/compile/test_async_tp.py
rename to tests/compile/distributed/test_async_tp.py
index 71ee228781438..86d409f1eadb0 100644
--- a/tests/compile/test_async_tp.py
+++ b/tests/compile/distributed/test_async_tp.py
@@ -27,13 +27,13 @@ from vllm.distributed.parallel_state import (
from vllm.platforms import current_platform
from vllm.utils.system_utils import update_environment_variables
-from ..models.registry import HF_EXAMPLE_MODELS
-from ..utils import (
+from ...models.registry import HF_EXAMPLE_MODELS
+from ...utils import (
compare_two_settings,
create_new_process_for_each_test,
multi_gpu_test,
)
-from .backend import TestBackend
+from ..backend import TestBackend
FP8_DTYPE = current_platform.fp8_dtype()
diff --git a/tests/compile/test_fusion_all_reduce.py b/tests/compile/distributed/test_fusion_all_reduce.py
similarity index 99%
rename from tests/compile/test_fusion_all_reduce.py
rename to tests/compile/distributed/test_fusion_all_reduce.py
index a539f4a160384..55133c4aaf4df 100644
--- a/tests/compile/test_fusion_all_reduce.py
+++ b/tests/compile/distributed/test_fusion_all_reduce.py
@@ -32,8 +32,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
from vllm.platforms import current_platform
from vllm.utils.system_utils import update_environment_variables
-from ..utils import TestFP8Layer, has_module_attribute, multi_gpu_test
-from .backend import TestBackend
+from ...utils import TestFP8Layer, has_module_attribute, multi_gpu_test
+from ..backend import TestBackend
class TestAllReduceRMSNormModel(torch.nn.Module):
diff --git a/tests/compile/test_fusions_e2e.py b/tests/compile/distributed/test_fusions_e2e.py
similarity index 98%
rename from tests/compile/test_fusions_e2e.py
rename to tests/compile/distributed/test_fusions_e2e.py
index f22d60ef000b2..661172e1965b5 100644
--- a/tests/compile/test_fusions_e2e.py
+++ b/tests/compile/distributed/test_fusions_e2e.py
@@ -18,7 +18,7 @@ from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer
from vllm.utils.torch_utils import is_torch_equal_or_newer
-from ..utils import flat_product, multi_gpu_test
+from ...utils import flat_product, multi_gpu_test
is_blackwell = lambda: current_platform.is_device_capability(100)
"""Are we running on Blackwell, a lot of tests depend on it"""
@@ -47,12 +47,8 @@ if current_platform.is_cuda():
ModelBackendTestCase(
# Use smaller model for L40s in CI
model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
- # TODO while llama4 is broken, use FLASHINFER for llama3 on Blackwell
- # so FI attention+fp8_quant is at least tested once
model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
- backend=AttentionBackendEnum.FLASHINFER
- if is_blackwell()
- else AttentionBackendEnum.TRITON_ATTN,
+ backend=AttentionBackendEnum.TRITON_ATTN,
matches=Matches(
attention_fusion=32,
allreduce_fusion=65,
@@ -65,9 +61,9 @@ if current_platform.is_cuda():
model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
# TODO FlashInfer attn broken on Hopper with kvcache=fp8:
# https://github.com/vllm-project/vllm/issues/28568
- # TODO FlashInfer attn broken on Blackwell for llama4:
- # https://github.com/vllm-project/vllm/issues/28604
- backend=AttentionBackendEnum.TRITON_ATTN,
+ backend=AttentionBackendEnum.FLASHINFER
+ if is_blackwell()
+ else AttentionBackendEnum.TRITON_ATTN,
matches=Matches(
attention_fusion=48,
allreduce_fusion=96,
diff --git a/tests/compile/test_sequence_parallelism.py b/tests/compile/distributed/test_sequence_parallelism.py
similarity index 99%
rename from tests/compile/test_sequence_parallelism.py
rename to tests/compile/distributed/test_sequence_parallelism.py
index d9675d93e2d14..1a0f623838f56 100644
--- a/tests/compile/test_sequence_parallelism.py
+++ b/tests/compile/distributed/test_sequence_parallelism.py
@@ -33,8 +33,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
from vllm.platforms import current_platform
from vllm.utils.system_utils import update_environment_variables
-from ..utils import TestFP8Layer, multi_gpu_test
-from .backend import TestBackend
+from ...utils import TestFP8Layer, multi_gpu_test
+from ..backend import TestBackend
FP8_DTYPE = current_platform.fp8_dtype()
prompts = [
diff --git a/tests/compile/fullgraph/__init__.py b/tests/compile/fullgraph/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/tests/compile/test_basic_correctness.py b/tests/compile/fullgraph/test_basic_correctness.py
similarity index 99%
rename from tests/compile/test_basic_correctness.py
rename to tests/compile/fullgraph/test_basic_correctness.py
index 3f6898607f6b9..965938c4433dd 100644
--- a/tests/compile/test_basic_correctness.py
+++ b/tests/compile/fullgraph/test_basic_correctness.py
@@ -7,7 +7,7 @@ import pytest
from vllm.config import CompilationMode
from vllm.utils.torch_utils import cuda_device_count_stateless
-from ..utils import compare_all_settings
+from ...utils import compare_all_settings
@dataclasses.dataclass
diff --git a/tests/compile/piecewise/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py
similarity index 100%
rename from tests/compile/piecewise/test_full_cudagraph.py
rename to tests/compile/fullgraph/test_full_cudagraph.py
diff --git a/tests/compile/test_full_graph.py b/tests/compile/fullgraph/test_full_graph.py
similarity index 99%
rename from tests/compile/test_full_graph.py
rename to tests/compile/fullgraph/test_full_graph.py
index b4e5e56ac9fe6..2c11ecef7f029 100644
--- a/tests/compile/test_full_graph.py
+++ b/tests/compile/fullgraph/test_full_graph.py
@@ -15,7 +15,7 @@ from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassC
from vllm.platforms import current_platform
from vllm.utils.torch_utils import is_torch_equal_or_newer
-from ..utils import create_new_process_for_each_test
+from ...utils import create_new_process_for_each_test
def models_list(*, all: bool = True, keywords: list[str] | None = None):
diff --git a/tests/compile/test_multimodal_compile.py b/tests/compile/fullgraph/test_multimodal_compile.py
similarity index 100%
rename from tests/compile/test_multimodal_compile.py
rename to tests/compile/fullgraph/test_multimodal_compile.py
diff --git a/tests/compile/piecewise/test_multiple_graphs.py b/tests/compile/fullgraph/test_multiple_graphs.py
similarity index 100%
rename from tests/compile/piecewise/test_multiple_graphs.py
rename to tests/compile/fullgraph/test_multiple_graphs.py
diff --git a/tests/compile/piecewise/test_simple.py b/tests/compile/fullgraph/test_simple.py
similarity index 100%
rename from tests/compile/piecewise/test_simple.py
rename to tests/compile/fullgraph/test_simple.py
diff --git a/tests/compile/piecewise/test_toy_llama.py b/tests/compile/fullgraph/test_toy_llama.py
similarity index 100%
rename from tests/compile/piecewise/test_toy_llama.py
rename to tests/compile/fullgraph/test_toy_llama.py
diff --git a/tests/compile/test_functionalization.py b/tests/compile/test_functionalization.py
index ef8ad92a923e6..0175029d72934 100644
--- a/tests/compile/test_functionalization.py
+++ b/tests/compile/test_functionalization.py
@@ -147,7 +147,7 @@ class TestRotaryEmbedding(torch.nn.Module):
self.head_dim,
rotary_dim=self.rotary_dim,
max_position=max_position,
- base=base,
+ rope_parameters={"rope_type": "default", "rope_theta": base},
)
def forward(self, positions, q, k):
@@ -182,7 +182,7 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=base,
+ rope_parameters={"rope_type": "default", "rope_theta": base},
)
def forward(self, positions, hidden_states):
diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py
new file mode 100644
index 0000000000000..1277c7e64eb21
--- /dev/null
+++ b/tests/config/test_config_utils.py
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from enum import Enum
+
+import pytest
+
+from vllm.config.utils import get_hash_factors, hash_factors, normalize_value
+
+# Helpers
+
+
+def endswith_fqname(obj, suffix: str) -> bool:
+ # normalize_value(type) returns fully-qualified name
+ # Compare suffix to avoid brittle import paths.
+ out = normalize_value(obj)
+ return isinstance(out, str) and out.endswith(suffix)
+
+
+def expected_path(p_str: str = ".") -> str:
+ import pathlib
+
+ p = pathlib.Path(p_str)
+ return p.expanduser().resolve().as_posix()
+
+
+# Minimal dataclass to test get_hash_factors.
+# Avoid importing heavy vLLM configs.
+@dataclass
+class SimpleConfig:
+ a: object
+ b: object | None = None
+
+
+class DummyLogprobsMode(Enum):
+ RAW_LOGITS = "raw_logits"
+
+
+def test_hash_factors_deterministic():
+ """Test that hash_factors produces consistent SHA-256 hashes"""
+ factors = {"a": 1, "b": "test"}
+ hash1 = hash_factors(factors)
+ hash2 = hash_factors(factors)
+
+ assert hash1 == hash2
+ # Dict key insertion order should not affect the hash.
+ factors_reordered = {"b": "test", "a": 1}
+ assert hash_factors(factors_reordered) == hash1
+ assert len(hash1) == 64
+ assert all(c in "0123456789abcdef" for c in hash1)
+
+
+@pytest.mark.parametrize(
+ "inp, expected",
+ [
+ (None, None),
+ (True, True),
+ (1, 1),
+ (1.0, 1.0),
+ ("x", "x"),
+ (b"ab", "6162"),
+ (bytearray(b"ab"), "6162"),
+ ([1, 2], (1, 2)),
+ ({"b": 2, "a": 1}, (("a", 1), ("b", 2))),
+ ],
+)
+def test_normalize_value_matrix(inp, expected):
+ """Parametric input→expected normalization table."""
+ assert normalize_value(inp) == expected
+
+
+def test_normalize_value_enum():
+ # Enums normalize to (module.QualName, value).
+ # DummyLogprobsMode uses a string payload.
+ out = normalize_value(DummyLogprobsMode.RAW_LOGITS)
+ assert isinstance(out, tuple)
+ assert out[0].endswith("DummyLogprobsMode")
+ # Expect string payload 'raw_logits'.
+ assert out[1] == "raw_logits"
+
+
+def test_normalize_value_set_order_insensitive():
+ # Sets are unordered; normalize_value sorts elements for determinism.
+ assert normalize_value({3, 1, 2}) == normalize_value({1, 2, 3})
+
+
+def test_normalize_value_path_normalization():
+ from pathlib import Path # local import to avoid global dependency
+
+ # Paths expand/resolve to absolute strings.
+ # Stabilizes hashing across working dirs.
+ assert normalize_value(Path(".")) == expected_path(".")
+
+
+def test_normalize_value_uuid_and_to_json():
+ # Objects may normalize via uuid() or to_json_string().
+ class HasUUID:
+ def uuid(self):
+ return "test-uuid"
+
+ class ToJson:
+ def to_json_string(self):
+ return '{"x":1}'
+
+ assert normalize_value(HasUUID()) == "test-uuid"
+ assert normalize_value(ToJson()) == '{"x":1}'
+
+
+@pytest.mark.parametrize(
+ "bad",
+ [
+ (lambda x: x),
+ (type("CallableInstance", (), {"__call__": lambda self: 0}))(),
+ (lambda: (lambda: 0))(), # nested function instance
+ ],
+)
+def test_error_cases(bad):
+ """Inputs expected to raise TypeError."""
+ # Reject functions/lambdas/callable instances
+ # to avoid under-hashing.
+ with pytest.raises(TypeError):
+ normalize_value(bad)
+
+
+def test_enum_vs_int_disambiguation():
+ # int stays primitive
+ nf_int = normalize_value(1)
+ assert nf_int == 1
+
+ # enum becomes ("module.QualName", value)
+ nf_enum = normalize_value(DummyLogprobsMode.RAW_LOGITS)
+ assert isinstance(nf_enum, tuple) and len(nf_enum) == 2
+ enum_type, enum_val = nf_enum
+ assert enum_type.endswith(".DummyLogprobsMode")
+ assert enum_val == "raw_logits"
+
+ # Build factor dicts from configs with int vs enum
+ f_int = get_hash_factors(SimpleConfig(1), set())
+ f_enum = get_hash_factors(SimpleConfig(DummyLogprobsMode.RAW_LOGITS), set())
+ # The int case remains a primitive value
+ assert f_int["a"] == 1
+ # The enum case becomes a tagged tuple ("module.QualName", "raw_logits")
+ assert isinstance(f_enum["a"], tuple) and f_enum["a"][1] == "raw_logits"
+ # Factor dicts must differ so we don't collide primitives with Enums.
+ assert f_int != f_enum
+ # Hash digests must differ correspondingly
+ assert hash_factors(f_int) != hash_factors(f_enum)
+
+ # Hash functions produce stable hex strings
+ h_int = hash_factors(f_int)
+ h_enum = hash_factors(f_enum)
+ assert isinstance(h_int, str) and len(h_int) == 64
+ assert isinstance(h_enum, str) and len(h_enum) == 64
+
+
+def test_classes_are_types():
+ """Types normalize to FQNs; include real vLLM types."""
+ # Only classes allowed; functions/lambdas are rejected.
+ # Canonical form is the fully-qualified name.
+ assert isinstance(normalize_value(str), str)
+
+ class LocalDummy:
+ pass
+
+ assert endswith_fqname(LocalDummy, ".LocalDummy")
diff --git a/tests/conftest.py b/tests/conftest.py
index b17081352edcf..163593eb3f14f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -748,6 +748,14 @@ class VllmRunner:
# being captured which can trigger edge cases that we don't handle yet.
kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
+ # Make sure we have atleast one cudagraph large enough for a single decode.
+ if (speculative_config := kwargs.get("speculative_config")) and (
+ num_speculative_tokens := speculative_config["num_speculative_tokens"]
+ ):
+ kwargs["compilation_config"]["cudagraph_capture_sizes"].append(
+ num_speculative_tokens + 1
+ )
+
with init_ctx:
self.llm = LLM(
model=model_name,
@@ -845,6 +853,7 @@ class VllmRunner:
@staticmethod
def _final_steps_generate_w_logprobs(
req_outputs: list[RequestOutput],
+ include_prompt_token_ids: bool = False,
) -> list[TokensTextLogprobsPromptLogprobs]:
outputs: list[TokensTextLogprobsPromptLogprobs] = []
for req_output in req_outputs:
@@ -853,9 +862,26 @@ class VllmRunner:
output_str = sample.text
output_ids = list(sample.token_ids)
output_logprobs = sample.logprobs
- outputs.append(
- (output_ids, output_str, output_logprobs, req_output.prompt_logprobs)
- )
+ if include_prompt_token_ids:
+ outputs.append(
+ ( # type: ignore[arg-type]
+ output_ids,
+ output_str,
+ output_logprobs,
+ req_output.prompt_token_ids,
+ req_output.prompt_logprobs,
+ )
+ )
+ else:
+ outputs.append(
+ (
+ output_ids,
+ output_str,
+ output_logprobs,
+ req_output.prompt_logprobs,
+ )
+ )
+
return outputs
def generate_w_logprobs(
@@ -865,6 +891,7 @@ class VllmRunner:
images: PromptImageInput | None = None,
audios: PromptAudioInput | None = None,
videos: PromptVideoInput | None = None,
+ include_prompt_token_ids: bool = False,
**kwargs: Any,
) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
@@ -874,7 +901,7 @@ class VllmRunner:
)
toks_str_logsprobs_prompt_logprobs = self._final_steps_generate_w_logprobs(
- req_outputs
+ req_outputs, include_prompt_token_ids
)
# Omit prompt logprobs if not required by sampling params
return (
diff --git a/tests/distributed/eplb_utils.py b/tests/distributed/eplb_utils.py
new file mode 100644
index 0000000000000..27a63e0215148
--- /dev/null
+++ b/tests/distributed/eplb_utils.py
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import random
+
+import torch
+import torch.multiprocessing as mp
+
+from vllm.distributed.parallel_state import (
+ init_distributed_environment,
+)
+from vllm.utils.system_utils import update_environment_variables
+
+mp.set_start_method("spawn", force=True)
+
+
+def distributed_run(fn, world_size, *args):
+ number_of_processes = world_size
+ processes: list[mp.Process] = []
+ for i in range(number_of_processes):
+ env: dict[str, str] = {}
+ env["RANK"] = str(i)
+ env["LOCAL_RANK"] = str(i)
+ env["WORLD_SIZE"] = str(number_of_processes)
+ env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
+ env["MASTER_ADDR"] = "localhost"
+ env["MASTER_PORT"] = "12345"
+ p = mp.Process(target=fn, args=(env, world_size, *args))
+ processes.append(p)
+ p.start()
+
+ for p in processes:
+ p.join()
+
+ for p in processes:
+ assert p.exitcode == 0
+
+
+def set_env_vars_and_device(env: dict[str, str]) -> None:
+ update_environment_variables(env)
+ local_rank = os.environ["LOCAL_RANK"]
+ device = torch.device(f"cuda:{local_rank}")
+ torch.cuda.set_device(device)
+ init_distributed_environment()
+
+ # Ensure each worker process has the same random seed
+ random.seed(42)
+ torch.manual_seed(42)
diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py
index b16fd0d06b145..7e4713b8aece0 100644
--- a/tests/distributed/test_context_parallel.py
+++ b/tests/distributed/test_context_parallel.py
@@ -31,7 +31,7 @@ class ParallelSetup(NamedTuple):
tp_size: int
pp_size: int
dcp_size: int
- dcp_kv_cache_interleave_size: int
+ cp_kv_cache_interleave_size: int
eager_mode: bool
chunked_prefill: bool
@@ -55,7 +55,7 @@ class CPTestSettings:
tp_base: int = 4,
pp_base: int = 1,
dcp_base: int = 1,
- dcp_kv_cache_interleave_size: int = 1,
+ cp_kv_cache_interleave_size: int = 1,
multi_node_only: bool = False,
runner: RunnerOption = "auto",
load_format: str | None = None,
@@ -71,7 +71,7 @@ class CPTestSettings:
tp_size=tp_base,
pp_size=pp_multiplier * pp_base,
dcp_size=int(dcp_multiplier * tp_base),
- dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
eager_mode=eager_mode_val,
chunked_prefill=chunked_prefill_val,
)
@@ -116,7 +116,7 @@ def _compare_cp_with_tp(
tp_size,
pp_size,
dcp_size,
- dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size,
eager_mode,
chunked_prefill,
) = parallel_setup
@@ -197,7 +197,7 @@ def _compare_cp_with_tp(
"--decode-context-parallel-size",
str(dcp_size),
"--dcp-kv-cache-interleave-size",
- str(dcp_kv_cache_interleave_size),
+ str(cp_kv_cache_interleave_size),
"--distributed-executor-backend",
distributed_backend,
]
@@ -227,7 +227,7 @@ CP_TEXT_GENERATION_MODELS = {
"deepseek-ai/DeepSeek-V2-Lite-Chat": [
CPTestSettings.detailed(),
CPTestSettings.detailed(tp_base=2),
- CPTestSettings.detailed(tp_base=2, dcp_kv_cache_interleave_size=64),
+ CPTestSettings.detailed(tp_base=2, cp_kv_cache_interleave_size=64),
],
"bigcode/gpt_bigcode-santacoder": [
CPTestSettings.detailed(),
diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py
index 7b45ae82c72d4..9498e75b279b7 100644
--- a/tests/distributed/test_eplb_execute.py
+++ b/tests/distributed/test_eplb_execute.py
@@ -1,8 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import multiprocessing
-import os
import random
import pytest
@@ -13,51 +11,9 @@ from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inp
from vllm.distributed.parallel_state import (
ensure_model_parallel_initialized,
get_tp_group,
- init_distributed_environment,
)
-from vllm.utils.system_utils import update_environment_variables
-
-def distributed_run(fn, world_size):
- number_of_processes = world_size
- processes: list[multiprocessing.Process] = []
- for i in range(number_of_processes):
- env: dict[str, str] = {}
- env["RANK"] = str(i)
- env["LOCAL_RANK"] = str(i)
- env["WORLD_SIZE"] = str(number_of_processes)
- env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
- env["MASTER_ADDR"] = "localhost"
- env["MASTER_PORT"] = "12345"
- p = multiprocessing.Process(target=fn, args=(env,))
- processes.append(p)
- p.start()
-
- for p in processes:
- p.join()
-
- for p in processes:
- assert p.exitcode == 0
-
-
-def worker_fn_wrapper(fn):
- # `multiprocessing.Process` cannot accept environment variables directly
- # so we need to pass the environment variables as arguments
- # and update the environment variables in the function
- def wrapped_fn(env):
- update_environment_variables(env)
- local_rank = os.environ["LOCAL_RANK"]
- device = torch.device(f"cuda:{local_rank}")
- torch.cuda.set_device(device)
- init_distributed_environment()
-
- # Ensure each worker process has the same random seed
- random.seed(42)
- torch.manual_seed(42)
-
- fn()
-
- return wrapped_fn
+from .eplb_utils import distributed_run, set_env_vars_and_device
def create_expert_indices_with_redundancy(
@@ -275,6 +231,79 @@ def verify_redundant_experts_have_same_weights(
)
+def _test_rearrange_expert_weights_with_redundancy(
+ env, world_size, num_layers, num_local_experts, num_logical_experts
+) -> None:
+ # Initialize model parallel (using tensor parallel as an entrypoint
+ # to expert parallel)
+ set_env_vars_and_device(env)
+ ensure_model_parallel_initialized(
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
+ )
+
+ ep_group = get_tp_group().cpu_group
+ ep_rank = torch.distributed.get_rank()
+ device = torch.device(f"cuda:{ep_rank}")
+
+ # Test parameters
+ total_physical_experts = world_size * num_local_experts
+ hidden_sizes = [32, 64] # Two different weight matrices
+
+ # Create old expert indices (with redundancy)
+ redundancy_config = create_redundancy_config(
+ num_logical_experts, total_physical_experts
+ )
+
+ old_indices = create_expert_indices_with_redundancy(
+ num_layers,
+ num_logical_experts,
+ total_physical_experts,
+ redundancy_config,
+ )
+
+ # Create new expert indices (with redundancy)
+ new_redundancy_config = create_redundancy_config(
+ num_logical_experts, total_physical_experts
+ )
+ new_indices = create_expert_indices_with_redundancy(
+ num_layers,
+ num_logical_experts,
+ total_physical_experts,
+ new_redundancy_config,
+ )
+
+ # Create expert weights
+ expert_weights = create_expert_weights(
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
+ )
+
+ # Execute weight rearrangement
+ rearrange_expert_weights_inplace(
+ old_indices,
+ new_indices,
+ expert_weights,
+ ep_group,
+ is_profile=False,
+ )
+
+ # Verify the rearrangement result
+ verify_expert_weights_after_shuffle(
+ expert_weights,
+ new_indices,
+ hidden_sizes,
+ ep_rank,
+ num_local_experts,
+ )
+
+ verify_redundant_experts_have_same_weights(
+ expert_weights,
+ new_indices,
+ hidden_sizes,
+ world_size,
+ num_local_experts,
+ )
+
+
@pytest.mark.parametrize(
"world_size,num_layers,num_local_experts,num_logical_experts",
[
@@ -305,78 +334,69 @@ def test_rearrange_expert_weights_with_redundancy(
if torch.cuda.device_count() < world_size:
pytest.skip(f"Need at least {world_size} GPUs to run the test")
+ distributed_run(
+ _test_rearrange_expert_weights_with_redundancy,
+ world_size,
+ num_layers,
+ num_local_experts,
+ num_logical_experts,
+ )
- @worker_fn_wrapper
- def worker_fn():
- # Initialize model parallel (using tensor parallel as an entrypoint
- # to expert parallel)
- ensure_model_parallel_initialized(
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
- )
- ep_group = get_tp_group().cpu_group
- ep_rank = torch.distributed.get_rank()
- device = torch.device(f"cuda:{ep_rank}")
+def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
+ set_env_vars_and_device(env)
+ ensure_model_parallel_initialized(
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
+ )
- # Test parameters
- total_physical_experts = world_size * num_local_experts
- hidden_sizes = [32, 64] # Two different weight matrices
+ ep_group = get_tp_group().cpu_group
+ ep_rank = torch.distributed.get_rank()
+ device = torch.device(f"cuda:{ep_rank}")
- # Create old expert indices (with redundancy)
- redundancy_config = create_redundancy_config(
- num_logical_experts, total_physical_experts
- )
+ num_layers = 2
+ num_local_experts = 2
+ total_physical_experts = world_size * num_local_experts
+ num_logical_experts = total_physical_experts // 2 # Some redundancy
+ hidden_sizes = [32, 64]
- old_indices = create_expert_indices_with_redundancy(
- num_layers,
- num_logical_experts,
- total_physical_experts,
- redundancy_config,
- )
+ # Create redundancy configuration
+ redundancy_config = [2] * num_logical_experts
- # Create new expert indices (with redundancy)
- new_redundancy_config = create_redundancy_config(
- num_logical_experts, total_physical_experts
- )
- new_indices = create_expert_indices_with_redundancy(
- num_layers,
- num_logical_experts,
- total_physical_experts,
- new_redundancy_config,
- )
+ # Same indices - no change
+ indices = create_expert_indices_with_redundancy(
+ num_layers, num_logical_experts, total_physical_experts, redundancy_config
+ )
- # Create expert weights
- expert_weights = create_expert_weights(
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
- )
+ expert_weights = create_expert_weights(
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
+ )
- # Execute weight rearrangement
- rearrange_expert_weights_inplace(
- old_indices,
- new_indices,
- expert_weights,
- ep_group,
- is_profile=False,
- )
+ # Save original weights
+ original_weights = []
+ for layer_weights in expert_weights:
+ layer_copy = []
+ for weight in layer_weights:
+ layer_copy.append(weight.clone())
+ original_weights.append(layer_copy)
- # Verify the rearrangement result
- verify_expert_weights_after_shuffle(
- expert_weights,
- new_indices,
- hidden_sizes,
- ep_rank,
- num_local_experts,
- )
+ # Execute rearrangement (should be no change)
+ rearrange_expert_weights_inplace(
+ indices,
+ indices, # Same indices
+ expert_weights,
+ ep_group,
+ is_profile=False,
+ )
- verify_redundant_experts_have_same_weights(
- expert_weights,
- new_indices,
- hidden_sizes,
- world_size,
- num_local_experts,
- )
-
- distributed_run(worker_fn, world_size)
+ # Verify that the weights have not changed
+ for layer in range(num_layers):
+ for weight_idx in range(len(hidden_sizes)):
+ torch.testing.assert_close(
+ expert_weights[layer][weight_idx],
+ original_weights[layer][weight_idx],
+ msg=f"""Layer {layer}, weight {weight_idx}
+ should remain unchanged""",
+ )
@pytest.mark.parametrize("world_size", [2, 4])
@@ -388,62 +408,69 @@ def test_rearrange_expert_weights_no_change(world_size):
if torch.cuda.device_count() < world_size:
pytest.skip(f"Need at least {world_size} GPUs to run the test")
+ distributed_run(_test_rearrange_expert_weights_no_change, world_size)
- @worker_fn_wrapper
- def worker_fn():
- ensure_model_parallel_initialized(
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
- )
- ep_group = get_tp_group().cpu_group
- ep_rank = torch.distributed.get_rank()
- device = torch.device(f"cuda:{ep_rank}")
+def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
+ set_env_vars_and_device(env)
+ ensure_model_parallel_initialized(
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
+ )
- num_layers = 2
- num_local_experts = 2
- total_physical_experts = world_size * num_local_experts
- num_logical_experts = total_physical_experts // 2 # Some redundancy
- hidden_sizes = [32, 64]
+ ep_group = get_tp_group().cpu_group
+ ep_rank = torch.distributed.get_rank()
+ device = torch.device(f"cuda:{ep_rank}")
- # Create redundancy configuration
- redundancy_config = [2] * num_logical_experts
+ num_layers = 1
+ num_local_experts = 2
+ total_physical_experts = world_size * num_local_experts
+ num_logical_experts = total_physical_experts // 2
+ hidden_sizes = [32]
- # Same indices - no change
- indices = create_expert_indices_with_redundancy(
- num_layers, num_logical_experts, total_physical_experts, redundancy_config
- )
+ # Create different index distributions
+ old_redundancy = create_redundancy_config(
+ num_logical_experts, total_physical_experts
+ )
+ new_redundancy = create_redundancy_config(
+ num_logical_experts, total_physical_experts
+ )
- expert_weights = create_expert_weights(
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
- )
+ old_indices = create_expert_indices_with_redundancy(
+ num_layers, num_logical_experts, total_physical_experts, old_redundancy
+ )
+ new_indices = create_expert_indices_with_redundancy(
+ num_layers, num_logical_experts, total_physical_experts, new_redundancy
+ )
- # Save original weights
- original_weights = []
- for layer_weights in expert_weights:
- layer_copy = []
- for weight in layer_weights:
- layer_copy.append(weight.clone())
- original_weights.append(layer_copy)
+ expert_weights = create_expert_weights(
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
+ )
- # Execute rearrangement (should be no change)
- rearrange_expert_weights_inplace(
- indices,
- indices, # Same indices
- expert_weights,
- ep_group,
- is_profile=False,
- )
+ # Save original weights
+ original_weights = []
+ for layer_weights in expert_weights:
+ layer_copy = []
+ for weight in layer_weights:
+ layer_copy.append(weight.clone())
+ original_weights.append(layer_copy)
- # Verify that the weights have not changed
- for layer in range(num_layers):
- for weight_idx in range(len(hidden_sizes)):
- torch.testing.assert_close(
- expert_weights[layer][weight_idx],
- original_weights[layer][weight_idx],
- msg=f"Layer {layer}, weight {weight_idx} should remain unchanged",
- )
+ # Execute profile mode rearrangement
+ rearrange_expert_weights_inplace(
+ old_indices,
+ new_indices,
+ expert_weights,
+ ep_group,
+ is_profile=True, # Profile mode
+ )
- distributed_run(worker_fn, world_size)
+ # In profile mode, the weights should remain unchanged
+ for layer in range(num_layers):
+ for weight_idx in range(len(hidden_sizes)):
+ torch.testing.assert_close(
+ expert_weights[layer][weight_idx],
+ original_weights[layer][weight_idx],
+ msg="In profile mode, the weights should remain unchanged",
+ )
@pytest.mark.parametrize("world_size", [2, 4])
@@ -452,66 +479,4 @@ def test_rearrange_expert_weights_profile_mode(world_size):
if torch.cuda.device_count() < world_size:
pytest.skip(f"Need at least {world_size} GPUs to run the test")
-
- @worker_fn_wrapper
- def worker_fn():
- ensure_model_parallel_initialized(
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
- )
-
- ep_group = get_tp_group().cpu_group
- ep_rank = torch.distributed.get_rank()
- device = torch.device(f"cuda:{ep_rank}")
-
- num_layers = 1
- num_local_experts = 2
- total_physical_experts = world_size * num_local_experts
- num_logical_experts = total_physical_experts // 2
- hidden_sizes = [32]
-
- # Create different index distributions
- old_redundancy = create_redundancy_config(
- num_logical_experts, total_physical_experts
- )
- new_redundancy = create_redundancy_config(
- num_logical_experts, total_physical_experts
- )
-
- old_indices = create_expert_indices_with_redundancy(
- num_layers, num_logical_experts, total_physical_experts, old_redundancy
- )
- new_indices = create_expert_indices_with_redundancy(
- num_layers, num_logical_experts, total_physical_experts, new_redundancy
- )
-
- expert_weights = create_expert_weights(
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
- )
-
- # Save original weights
- original_weights = []
- for layer_weights in expert_weights:
- layer_copy = []
- for weight in layer_weights:
- layer_copy.append(weight.clone())
- original_weights.append(layer_copy)
-
- # Execute profile mode rearrangement
- rearrange_expert_weights_inplace(
- old_indices,
- new_indices,
- expert_weights,
- ep_group,
- is_profile=True, # Profile mode
- )
-
- # In profile mode, the weights should remain unchanged
- for layer in range(num_layers):
- for weight_idx in range(len(hidden_sizes)):
- torch.testing.assert_close(
- expert_weights[layer][weight_idx],
- original_weights[layer][weight_idx],
- msg="In profile mode, the weights should remain unchanged",
- )
-
- distributed_run(worker_fn, world_size)
+ distributed_run(_test_rearrange_expert_weights_profile_mode, world_size)
diff --git a/tests/distributed/test_eplb_fused_moe_layer.py b/tests/distributed/test_eplb_fused_moe_layer.py
new file mode 100644
index 0000000000000..55f26519887a1
--- /dev/null
+++ b/tests/distributed/test_eplb_fused_moe_layer.py
@@ -0,0 +1,285 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Test that the interaction between EPLB and FusedMoE Layer is okay
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
+from vllm.distributed.parallel_state import (
+ ensure_model_parallel_initialized,
+ get_tp_group,
+)
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+
+from .eplb_utils import distributed_run, set_env_vars_and_device
+
+
+@dataclass
+class TestConfig:
+ num_layers: int
+ num_experts: int
+ num_local_experts: int
+ num_topk: int
+ hidden_size: int
+ intermediate_size: int
+ weight_dtype: torch.dtype
+ weight_scale_dtype: torch.dtype | None
+ column_major_scales: bool
+
+
+def make_expert_weights(
+ layer_idx: int,
+ global_expert_idx: int,
+ global_num_experts: int,
+ tensor_shape: tuple[int, ...],
+ tensor_dtype: torch.dtype,
+ tensor_device: torch.device,
+ is_column_major: bool,
+) -> torch.Tensor:
+ assert len(tensor_shape) == 2
+
+ if is_column_major:
+ tensor_shape = (tensor_shape[1], tensor_shape[0])
+
+ x = torch.empty(tensor_shape, dtype=tensor_dtype, device=tensor_device)
+ value_offset = (layer_idx * global_num_experts + global_expert_idx) * x.numel()
+ x.view(-1).copy_(
+ torch.arange(
+ value_offset,
+ value_offset + x.numel(),
+ dtype=tensor_dtype,
+ device=tensor_device,
+ )
+ )
+
+ if is_column_major:
+ x = torch.transpose(x, 1, 0)
+ assert not x.is_contiguous()
+ return x
+
+
+def make_fused_moe_layer(
+ rank: int,
+ layer_idx: int,
+ test_config: TestConfig,
+) -> FusedMoE:
+ fml = FusedMoE(
+ num_experts=test_config.num_experts,
+ top_k=test_config.num_topk,
+ hidden_size=test_config.hidden_size,
+ intermediate_size=test_config.intermediate_size,
+ prefix=f"dummy_layer_{layer_idx}",
+ activation="silu",
+ is_act_and_mul=True,
+ params_dtype=test_config.weight_dtype,
+ )
+
+ device = torch.device(f"cuda:{rank}")
+
+ from functools import partial
+
+ _make_expert_weights = partial(
+ make_expert_weights,
+ layer_idx=layer_idx,
+ global_num_experts=test_config.num_experts,
+ tensor_device=device,
+ )
+
+ assert isinstance(fml.w13_weight.data, torch.Tensor)
+ assert isinstance(fml.w2_weight.data, torch.Tensor)
+ fml.w13_weight.data = fml.w13_weight.data.to(device=device)
+ fml.w2_weight.data = fml.w2_weight.data.to(device=device)
+ w13_weight = fml.w13_weight.data
+ w2_weight = fml.w2_weight.data
+ assert w13_weight.size(0) == test_config.num_local_experts
+ for i in range(test_config.num_local_experts):
+ g_i = rank * test_config.num_local_experts + i
+ w13_weight_e = w13_weight[i]
+ w2_weight_e = w2_weight[i]
+ w13_weight_e.copy_(
+ _make_expert_weights(
+ global_expert_idx=g_i,
+ tensor_shape=w13_weight_e.shape,
+ tensor_dtype=w13_weight_e.dtype,
+ is_column_major=False,
+ )
+ )
+ w2_weight_e.copy_(
+ _make_expert_weights(
+ global_expert_idx=g_i,
+ tensor_shape=w2_weight_e.shape,
+ tensor_dtype=w2_weight_e.dtype,
+ is_column_major=False,
+ )
+ )
+
+ block_size = 16
+
+ def block_quant_scales_shape(
+ shape: tuple[int, ...], is_column_major: bool
+ ) -> tuple[int, ...]:
+ assert len(shape) == 3
+ if not is_column_major:
+ return (shape[0], shape[1] // block_size, shape[2] // block_size)
+ else:
+ return (shape[0], shape[2] // block_size, shape[1] // block_size)
+
+ is_column_major = test_config.column_major_scales
+ w13_weight_scale_inv = torch.empty(
+ block_quant_scales_shape(w13_weight.shape, is_column_major),
+ dtype=test_config.weight_dtype,
+ device=device,
+ )
+ w2_weight_scale_inv = torch.empty(
+ block_quant_scales_shape(w2_weight.shape, is_column_major),
+ dtype=test_config.weight_dtype,
+ device=device,
+ )
+
+ for i in range(test_config.num_local_experts):
+ g_i = rank * test_config.num_local_experts + i
+ w13_s_e = w13_weight_scale_inv[i]
+ w2_s_e = w2_weight_scale_inv[i]
+ w13_s_e.copy_(
+ _make_expert_weights(
+ global_expert_idx=g_i,
+ tensor_shape=w13_s_e.shape,
+ tensor_dtype=w13_s_e.dtype,
+ # Fill data in row-major and then
+ # transpose if test_config requires col-major.
+ is_column_major=False,
+ )
+ )
+ w2_s_e.copy_(
+ _make_expert_weights(
+ global_expert_idx=g_i,
+ tensor_shape=w2_s_e.shape,
+ tensor_dtype=w2_s_e.dtype,
+ is_column_major=False,
+ )
+ )
+ if is_column_major:
+ w13_weight_scale_inv = torch.transpose(w13_weight_scale_inv, 1, 2)
+ w2_weight_scale_inv = torch.transpose(w2_weight_scale_inv, 1, 2)
+ assert not w13_weight_scale_inv.is_contiguous()
+ assert not w2_weight_scale_inv.is_contiguous()
+
+ # Add scales to the parameter list
+ fml.w13_weight_scale_inv = torch.nn.Parameter(
+ w13_weight_scale_inv, requires_grad=False
+ )
+ fml.w2_weight_scale_inv = torch.nn.Parameter(
+ w2_weight_scale_inv, requires_grad=False
+ )
+
+ return fml
+
+
+def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
+ # Initialize model parallel (using tensor parallel as an entrypoint
+ # to expert parallel)
+ set_env_vars_and_device(env)
+
+ vllm_config = VllmConfig()
+ vllm_config.parallel_config.tensor_parallel_size = world_size
+ vllm_config.parallel_config.enable_expert_parallel = True
+
+ with set_current_vllm_config(vllm_config):
+ ensure_model_parallel_initialized(
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
+ )
+
+ ep_group = get_tp_group().cpu_group
+ ep_rank = torch.distributed.get_rank()
+
+ fml_layers = [
+ make_fused_moe_layer(ep_rank, layer_idx, test_config)
+ for layer_idx in range(test_config.num_layers)
+ ]
+ rank_expert_weights = [fml.get_expert_weights() for fml in fml_layers]
+
+ indices = torch.zeros(
+ test_config.num_layers, test_config.num_experts, dtype=torch.long
+ )
+ for lidx in range(test_config.num_layers):
+ indices[lidx] = torch.Tensor(range(test_config.num_experts))
+
+ shuffled_indices = torch.zeros_like(indices)
+ for lidx in range(test_config.num_layers):
+ shuffled_indices[lidx] = torch.randperm(test_config.num_experts)
+
+ rearrange_expert_weights_inplace(
+ indices,
+ shuffled_indices,
+ rank_expert_weights,
+ ep_group,
+ is_profile=False,
+ )
+
+ num_local_experts = test_config.num_local_experts
+ num_global_experts = test_config.num_experts
+ for lidx, fml in enumerate(fml_layers):
+ for name, w in fml.named_parameters():
+ for e in range(num_local_experts):
+ g_e = shuffled_indices[lidx][ep_rank * num_local_experts + e]
+ ref = make_expert_weights(
+ layer_idx=lidx,
+ global_expert_idx=int(g_e.item()),
+ global_num_experts=num_global_experts,
+ tensor_shape=w[e].shape,
+ tensor_dtype=w[e].dtype,
+ tensor_device=w[e].device,
+ is_column_major=not w[e].is_contiguous(),
+ )
+ assert w[e].shape == ref.shape and w[e].stride() == ref.stride(), (
+ f"w[{e}] {w[e].size()} {w[e].stride()} vs "
+ f"ref {ref.size()} {ref.stride()}"
+ )
+ torch.testing.assert_close(w[e], ref)
+
+
+@pytest.mark.parametrize("world_size", [2])
+@pytest.mark.parametrize("num_layers", [4])
+@pytest.mark.parametrize("num_experts", [16])
+@pytest.mark.parametrize("hidden_size", [256])
+@pytest.mark.parametrize("intermediate_size", [256])
+@pytest.mark.parametrize("column_major_scales", [True, False])
+def test_eplb_fml(
+ world_size: int,
+ num_layers: int,
+ num_experts: int,
+ hidden_size: int,
+ intermediate_size: int,
+ column_major_scales: bool,
+):
+ if torch.cuda.device_count() < world_size:
+ pytest.skip(f"Need at least {world_size} GPUs to run the test")
+
+ num_local_experts = num_experts // world_size
+ num_topk = 4
+ # The dtypes are fine as we are essentially just checking data-copies
+ weight_dtype = torch.bfloat16
+ weight_scale_dtype = torch.bfloat16
+
+ test_config = TestConfig(
+ num_layers=num_layers,
+ num_experts=num_experts,
+ num_local_experts=num_local_experts,
+ num_topk=num_topk,
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ weight_dtype=weight_dtype,
+ weight_scale_dtype=weight_scale_dtype,
+ column_major_scales=column_major_scales,
+ )
+
+ distributed_run(
+ _test_eplb_fml,
+ world_size,
+ test_config,
+ )
diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py
index 0ab94d30858fb..89f035d2cdd6f 100644
--- a/tests/distributed/test_pipeline_parallel.py
+++ b/tests/distributed/test_pipeline_parallel.py
@@ -130,6 +130,7 @@ TEXT_GENERATION_MODELS = {
"inceptionai/jais-13b-chat": PPTestSettings.fast(),
"ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
"pfnet/plamo-2-1b": PPTestSettings.fast(),
+ "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(),
"meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(),
# Tests TransformersForCausalLM
"hmellor/Ilama-3.2-1B": PPTestSettings.fast(),
diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py
index c3085beeb3564..c7c9d0602def0 100644
--- a/tests/distributed/test_pynccl.py
+++ b/tests/distributed/test_pynccl.py
@@ -1,9 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import multiprocessing
import os
+import multiprocess as mp
import numpy as np
import pytest
import torch
@@ -20,10 +20,12 @@ from vllm.distributed.parallel_state import (
)
from vllm.utils.system_utils import update_environment_variables
+mp.set_start_method("spawn", force=True)
+
def distributed_run(fn, world_size):
number_of_processes = world_size
- processes: list[multiprocessing.Process] = []
+ processes: list[mp.Process] = []
for i in range(number_of_processes):
env: dict[str, str] = {}
env["RANK"] = str(i)
@@ -32,7 +34,7 @@ def distributed_run(fn, world_size):
env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
env["MASTER_ADDR"] = "localhost"
env["MASTER_PORT"] = "12345"
- p = multiprocessing.Process(target=fn, args=(env,))
+ p = mp.Process(target=fn, args=(env,))
processes.append(p)
p.start()
diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py
index 472b1487ef440..10827e3b4b9cd 100644
--- a/tests/engine/test_arg_utils.py
+++ b/tests/engine/test_arg_utils.py
@@ -279,7 +279,7 @@ def test_prefix_cache_default():
args = parser.parse_args([])
engine_args = EngineArgs.from_cli_args(args=args)
- assert not engine_args.enable_prefix_caching, "prefix caching defaults to off."
+ assert engine_args.enable_prefix_caching, "prefix caching should default to on."
# with flag to turn it on.
args = parser.parse_args(["--enable-prefix-caching"])
diff --git a/tests/entrypoints/openai/test_enable_force_include_usage.py b/tests/entrypoints/openai/test_enable_force_include_usage.py
index 3ddf2308eb1d5..9d527c45c1fae 100644
--- a/tests/entrypoints/openai/test_enable_force_include_usage.py
+++ b/tests/entrypoints/openai/test_enable_force_include_usage.py
@@ -17,7 +17,7 @@ def chat_server_with_force_include_usage(request): # noqa: F811
"128",
"--enforce-eager",
"--max-num-seqs",
- "1",
+ "4",
"--enable-force-include-usage",
"--port",
"55857",
@@ -78,7 +78,7 @@ def transcription_server_with_force_include_usage():
"--dtype",
"bfloat16",
"--max-num-seqs",
- "1",
+ "4",
"--enforce-eager",
"--enable-force-include-usage",
"--gpu-memory-utilization",
diff --git a/tests/entrypoints/openai/test_metrics.py b/tests/entrypoints/openai/test_metrics.py
index dbcec9d31fc9b..4e7b765d7713f 100644
--- a/tests/entrypoints/openai/test_metrics.py
+++ b/tests/entrypoints/openai/test_metrics.py
@@ -16,6 +16,7 @@ from transformers import AutoTokenizer
from vllm import version
+from ...conftest import LocalAssetServer
from ...utils import RemoteOpenAIServer
MODELS = {
@@ -69,7 +70,6 @@ async def client(server):
_PROMPT = "Hello my name is Robert and I love magic"
-_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
def _get_expected_values(num_requests: int, prompt_ids: list[int], max_tokens: int):
@@ -250,6 +250,7 @@ HIDDEN_DEPRECATED_METRICS: list[str] = [
@pytest.mark.asyncio
async def test_metrics_exist(
+ local_asset_server: LocalAssetServer,
server: RemoteOpenAIServer,
client: openai.AsyncClient,
model_key: str,
@@ -265,13 +266,21 @@ async def test_metrics_exist(
temperature=0.0,
)
else:
+ # https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
await client.chat.completions.create(
model=model_name,
messages=[
{
"role": "user",
"content": [
- {"type": "image_url", "image_url": {"url": _IMAGE_URL}},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": local_asset_server.url_for(
+ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ ),
+ },
+ },
{"type": "text", "text": "What's in this image?"},
],
}
diff --git a/tests/entrypoints/openai/test_response_api_simple.py b/tests/entrypoints/openai/test_response_api_simple.py
new file mode 100644
index 0000000000000..425b8199a0fd0
--- /dev/null
+++ b/tests/entrypoints/openai/test_response_api_simple.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import pytest
+import pytest_asyncio
+from openai import OpenAI
+
+from ...utils import RemoteOpenAIServer
+
+MODEL_NAME = "Qwen/Qwen3-8B"
+
+
+@pytest.fixture(scope="module")
+def server():
+ args = ["--reasoning-parser", "qwen3", "--max_model_len", "5000"]
+ env_dict = dict(
+ VLLM_ENABLE_RESPONSES_API_STORE="1",
+ # uncomment for tool calling
+ # PYTHON_EXECUTION_BACKEND="dangerously_use_uv",
+ )
+
+ with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
+ yield remote_server
+
+
+@pytest_asyncio.fixture
+async def client(server):
+ async with server.get_async_client() as async_client:
+ yield async_client
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_basic(client: OpenAI, model_name: str):
+ response = await client.responses.create(
+ model=model_name,
+ input="What is 13 * 24?",
+ )
+ assert response is not None
+ print("response: ", response)
+ assert response.status == "completed"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_reasoning_item(client: OpenAI, model_name: str):
+ response = await client.responses.create(
+ model=model_name,
+ input=[
+ {"type": "message", "content": "Hello.", "role": "user"},
+ {
+ "type": "reasoning",
+ "id": "lol",
+ "content": [
+ {
+ "type": "reasoning_text",
+ "text": "We need to respond: greeting.",
+ }
+ ],
+ "summary": [],
+ },
+ ],
+ temperature=0.0,
+ )
+ assert response is not None
+ assert response.status == "completed"
+ # make sure we get a reasoning and text output
+ assert response.output[0].type == "reasoning"
+ assert response.output[1].type == "message"
+ assert type(response.output[1].content[0].text) is str
diff --git a/tests/entrypoints/openai/test_response_api_with_harmony.py b/tests/entrypoints/openai/test_response_api_with_harmony.py
index dea8d2d28f61a..6251e1776c30a 100644
--- a/tests/entrypoints/openai/test_response_api_with_harmony.py
+++ b/tests/entrypoints/openai/test_response_api_with_harmony.py
@@ -35,7 +35,7 @@ GET_WEATHER_SCHEMA = {
@pytest.fixture(scope="module")
def server():
- args = ["--enforce-eager", "--tool-server", "demo"]
+ args = ["--enforce-eager", "--tool-server", "demo", "--max_model_len", "5000"]
env_dict = dict(
VLLM_ENABLE_RESPONSES_API_STORE="1",
PYTHON_EXECUTION_BACKEND="dangerously_use_uv",
@@ -550,6 +550,31 @@ def call_function(name, args):
raise ValueError(f"Unknown function: {name}")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_reasoning_item(client: OpenAI, model_name: str):
+ response = await client.responses.create(
+ model=model_name,
+ input=[
+ {"type": "message", "content": "Hello.", "role": "user"},
+ {
+ "type": "reasoning",
+ "id": "lol",
+ "content": [
+ {
+ "type": "reasoning_text",
+ "text": "We need to respond: greeting.",
+ }
+ ],
+ "summary": [],
+ },
+ ],
+ temperature=0.0,
+ )
+ assert response is not None
+ assert response.status == "completed"
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_function_calling(client: OpenAI, model_name: str):
diff --git a/tests/entrypoints/openai/test_vision.py b/tests/entrypoints/openai/test_vision.py
index 2a7df08ea3b0e..d83c6726e72da 100644
--- a/tests/entrypoints/openai/test_vision.py
+++ b/tests/entrypoints/openai/test_vision.py
@@ -17,10 +17,10 @@ MAXIMUM_IMAGES = 2
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
- "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
- "Grayscale_8bits_palette_sample_image.png", # "https://upload.wikimedia.org/wikipedia/commons/f/fa/Grayscale_8bits_palette_sample_image.png",
- "1280px-Venn_diagram_rgb.svg.png", # "https://upload.wikimedia.org/wikipedia/commons/thumb/9/91/Venn_diagram_rgb.svg/1280px-Venn_diagram_rgb.svg.png",
- "RGBA_comp.png", # "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
+ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
+ "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",
]
EXPECTED_MM_BEAM_SEARCH_RES = [
diff --git a/tests/entrypoints/pooling/correctness/test_mteb_embed.py b/tests/entrypoints/pooling/correctness/test_mteb_embed.py
index 7f16638e51e2c..64673534fd32a 100644
--- a/tests/entrypoints/pooling/correctness/test_mteb_embed.py
+++ b/tests/entrypoints/pooling/correctness/test_mteb_embed.py
@@ -11,6 +11,12 @@ from tests.models.language.pooling_mteb_test.mteb_utils import (
run_mteb_embed_task,
)
from tests.utils import RemoteOpenAIServer
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
diff --git a/tests/entrypoints/pooling/correctness/test_mteb_score.py b/tests/entrypoints/pooling/correctness/test_mteb_score.py
index 1afe68b189db8..81ad0097187b0 100644
--- a/tests/entrypoints/pooling/correctness/test_mteb_score.py
+++ b/tests/entrypoints/pooling/correctness/test_mteb_score.py
@@ -13,6 +13,12 @@ from tests.models.language.pooling_mteb_test.mteb_utils import (
run_mteb_rerank,
)
from tests.utils import RemoteOpenAIServer
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
diff --git a/tests/entrypoints/pooling/llm/test_embedding.py b/tests/entrypoints/pooling/llm/test_embedding.py
index 5455b5f91fc09..f5eab4c29ae18 100644
--- a/tests/entrypoints/pooling/llm/test_embedding.py
+++ b/tests/entrypoints/pooling/llm/test_embedding.py
@@ -9,6 +9,12 @@ import torch.nn.functional as F
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODEL_NAME = "intfloat/multilingual-e5-small"
diff --git a/tests/entrypoints/pooling/llm/test_encode.py b/tests/entrypoints/pooling/llm/test_encode.py
index ca85d2758fce4..f86ecef2e4744 100644
--- a/tests/entrypoints/pooling/llm/test_encode.py
+++ b/tests/entrypoints/pooling/llm/test_encode.py
@@ -7,6 +7,12 @@ import pytest
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODEL_NAME = "intfloat/multilingual-e5-small"
diff --git a/tests/entrypoints/pooling/llm/test_score.py b/tests/entrypoints/pooling/llm/test_score.py
index b69c6a47c1913..ce36d61cb8476 100644
--- a/tests/entrypoints/pooling/llm/test_score.py
+++ b/tests/entrypoints/pooling/llm/test_score.py
@@ -9,6 +9,12 @@ import torch
from tests.models.utils import softmax
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
diff --git a/tests/entrypoints/pooling/openai/test_embedding.py b/tests/entrypoints/pooling/openai/test_embedding.py
index e971b23e8f1a0..0c88d800e2f99 100644
--- a/tests/entrypoints/pooling/openai/test_embedding.py
+++ b/tests/entrypoints/pooling/openai/test_embedding.py
@@ -19,6 +19,7 @@ from vllm.entrypoints.openai.protocol import (
EmbeddingResponse,
PoolingResponse,
)
+from vllm.platforms import current_platform
from vllm.transformers_utils.tokenizer import get_tokenizer
from vllm.utils.serial_utils import (
EMBED_DTYPE_TO_TORCH_DTYPE,
@@ -28,6 +29,11 @@ from vllm.utils.serial_utils import (
decode_pooling_output,
)
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
+
MODEL_NAME = "intfloat/multilingual-e5-small"
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
DTYPE = "bfloat16"
diff --git a/tests/entrypoints/pooling/openai/test_embedding_dimensions.py b/tests/entrypoints/pooling/openai/test_embedding_dimensions.py
index ba9fb64262772..8018dac2d3ffe 100644
--- a/tests/entrypoints/pooling/openai/test_embedding_dimensions.py
+++ b/tests/entrypoints/pooling/openai/test_embedding_dimensions.py
@@ -12,6 +12,12 @@ from tests.models.language.pooling.embed_utils import run_embedding_correctness_
from tests.models.utils import EmbedModelInfo
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.openai.protocol import EmbeddingResponse
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODELS = [
EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
diff --git a/tests/entrypoints/pooling/openai/test_embedding_long_text.py b/tests/entrypoints/pooling/openai/test_embedding_long_text.py
index f977c81a9084e..a9ade09dad0b5 100644
--- a/tests/entrypoints/pooling/openai/test_embedding_long_text.py
+++ b/tests/entrypoints/pooling/openai/test_embedding_long_text.py
@@ -16,6 +16,12 @@ import pytest_asyncio
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.openai.protocol import EmbeddingResponse
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
def _generate_random_text(word_count: int) -> str:
diff --git a/tests/entrypoints/pooling/openai/test_rerank.py b/tests/entrypoints/pooling/openai/test_rerank.py
index 1d85190c12a19..5a772e22a7414 100644
--- a/tests/entrypoints/pooling/openai/test_rerank.py
+++ b/tests/entrypoints/pooling/openai/test_rerank.py
@@ -8,6 +8,12 @@ import torch.nn.functional as F
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.openai.protocol import PoolingResponse, RerankResponse
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODEL_NAME = "BAAI/bge-reranker-base"
DTYPE = "bfloat16"
diff --git a/tests/entrypoints/pooling/openai/test_score.py b/tests/entrypoints/pooling/openai/test_score.py
index b8f796d47efaa..ceff9d0181825 100644
--- a/tests/entrypoints/pooling/openai/test_score.py
+++ b/tests/entrypoints/pooling/openai/test_score.py
@@ -10,6 +10,12 @@ from torch import tensor
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.openai.protocol import ScoreResponse
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODELS = [
{"name": "BAAI/bge-reranker-v2-m3", "is_cross_encoder": True},
diff --git a/tests/entrypoints/pooling/openai/test_truncation.py b/tests/entrypoints/pooling/openai/test_truncation.py
index 6889628dc9145..0d2d385840402 100644
--- a/tests/entrypoints/pooling/openai/test_truncation.py
+++ b/tests/entrypoints/pooling/openai/test_truncation.py
@@ -7,6 +7,12 @@ import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+ pytest.skip(
+ "Encoder self-attention is not implemented on ROCm.", allow_module_level=True
+ )
MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
max_model_len = 128
diff --git a/tests/entrypoints/pooling/openai/test_vision_embedding.py b/tests/entrypoints/pooling/openai/test_vision_embedding.py
index 944392d66fa5f..1befb5a3cf7a8 100644
--- a/tests/entrypoints/pooling/openai/test_vision_embedding.py
+++ b/tests/entrypoints/pooling/openai/test_vision_embedding.py
@@ -19,10 +19,10 @@ assert vlm2vec_jinja_path.exists()
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
- "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
- "Grayscale_8bits_palette_sample_image.png", # "https://upload.wikimedia.org/wikipedia/commons/f/fa/Grayscale_8bits_palette_sample_image.png",
- "1280px-Venn_diagram_rgb.svg.png", # "https://upload.wikimedia.org/wikipedia/commons/thumb/9/91/Venn_diagram_rgb.svg/1280px-Venn_diagram_rgb.svg.png",
- "RGBA_comp.png", # "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
+ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
+ "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",
]
diff --git a/tests/entrypoints/test_chat_utils.py b/tests/entrypoints/test_chat_utils.py
index ca87b3e76b3f4..7baf564ad01a4 100644
--- a/tests/entrypoints/test_chat_utils.py
+++ b/tests/entrypoints/test_chat_utils.py
@@ -103,6 +103,19 @@ def qwen2_audio_model_config():
)
+@pytest.fixture(scope="function")
+def audio_embeds_model_config():
+ return ModelConfig(
+ QWEN2AUDIO_MODEL_ID,
+ runner="generate",
+ trust_remote_code=True,
+ limit_mm_per_prompt={
+ "audio": 2,
+ },
+ enable_mm_embeds=True,
+ )
+
+
@pytest.fixture(scope="module")
def qwen2_audio_tokenizer():
return get_tokenizer(QWEN2AUDIO_MODEL_ID)
@@ -843,6 +856,138 @@ def test_parse_chat_messages_empty_image_embeds_with_uuid(
_assert_mm_uuids(mm_uuids, 1, expected_uuids=[uuid])
+def test_parse_chat_messages_empty_audio_embeds_with_uuid(
+ audio_embeds_model_config,
+ qwen2_audio_tokenizer,
+):
+ """Test audio_embeds with UUID (no actual embeds data)."""
+ uuid = "test-audio-uuid-123"
+
+ conversation, mm_data, mm_uuids = parse_chat_messages(
+ [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Describe this audio"},
+ {"type": "audio_embeds", "audio_embeds": None, "uuid": uuid},
+ ],
+ }
+ ],
+ audio_embeds_model_config,
+ qwen2_audio_tokenizer,
+ content_format="string",
+ )
+
+ # Should have audio in mm_data as None (UUID provided)
+ assert mm_data is not None
+ assert "audio" in mm_data
+ assert mm_data["audio"] is None
+ # UUID should be recorded
+ assert mm_uuids is not None
+ assert "audio" in mm_uuids
+ _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[uuid])
+
+
+def test_parse_chat_messages_audio_embeds_with_string(
+ audio_embeds_model_config,
+ qwen2_audio_tokenizer,
+):
+ """Test audio_embeds with base64 string embedding data."""
+ import base64
+ import io
+
+ import torch
+
+ # Create a sample audio embedding tensor
+ audio_embedding = torch.randn(1, 128, 768)
+
+ # Encode it as base64
+ buffer = io.BytesIO()
+ torch.save(audio_embedding, buffer)
+ buffer.seek(0)
+ binary_data = buffer.read()
+ base64_audio_embedding = base64.b64encode(binary_data).decode("utf-8")
+
+ conversation, mm_data, mm_uuids = parse_chat_messages(
+ [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Describe this audio"},
+ {
+ "type": "audio_embeds",
+ "audio_embeds": base64_audio_embedding,
+ },
+ ],
+ }
+ ],
+ audio_embeds_model_config,
+ qwen2_audio_tokenizer,
+ content_format="string",
+ )
+
+ # Should have audio embedding in mm_data (single tensor, not a list)
+ assert mm_data is not None
+ assert "audio" in mm_data
+ assert isinstance(mm_data["audio"], torch.Tensor)
+ assert mm_data["audio"].shape == audio_embedding.shape
+ # No UUID provided
+ assert mm_uuids is not None
+ assert "audio" in mm_uuids
+ _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[None])
+
+
+@pytest.mark.asyncio
+async def test_parse_chat_messages_audio_embeds_async(
+ audio_embeds_model_config,
+ qwen2_audio_tokenizer,
+):
+ """Test audio_embeds with async futures."""
+ import base64
+ import io
+
+ import torch
+
+ # Create a sample audio embedding tensor
+ audio_embedding = torch.randn(1, 128, 768)
+
+ # Encode it as base64
+ buffer = io.BytesIO()
+ torch.save(audio_embedding, buffer)
+ buffer.seek(0)
+ binary_data = buffer.read()
+ base64_audio_embedding = base64.b64encode(binary_data).decode("utf-8")
+
+ conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+ [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Describe this audio"},
+ {
+ "type": "audio_embeds",
+ "audio_embeds": base64_audio_embedding,
+ },
+ ],
+ }
+ ],
+ audio_embeds_model_config,
+ qwen2_audio_tokenizer,
+ content_format="string",
+ )
+
+ # Should have audio embedding in mm_data (single tensor, not a list)
+ mm_data = await mm_future
+ assert mm_data is not None
+ assert "audio" in mm_data
+ assert isinstance(mm_data["audio"], torch.Tensor)
+ assert mm_data["audio"].shape == audio_embedding.shape
+ # No UUID provided
+ assert mm_uuids is not None
+ assert "audio" in mm_uuids
+ _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[None])
+
+
@pytest.mark.asyncio
async def test_parse_chat_messages_empty_image_embeds_with_uuid_async(
phi3v_model_config_image_embeds,
diff --git a/tests/entrypoints/test_responses_utils.py b/tests/entrypoints/test_responses_utils.py
index 48bf06088bc05..91c818374e3fd 100644
--- a/tests/entrypoints/test_responses_utils.py
+++ b/tests/entrypoints/test_responses_utils.py
@@ -1,7 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+from openai.types.responses.response_reasoning_item import (
+ Content,
+ ResponseReasoningItem,
+ Summary,
+)
+
from vllm.entrypoints.responses_utils import (
+ construct_chat_message_with_tool_call,
convert_tool_responses_to_completions_format,
)
@@ -28,3 +36,53 @@ class TestResponsesUtils:
result = convert_tool_responses_to_completions_format(input_tool)
assert result == {"type": "function", "function": input_tool}
+
+ def test_construct_chat_message_with_tool_call(self):
+ item = ResponseReasoningItem(
+ id="lol",
+ summary=[],
+ type="reasoning",
+ content=[
+ Content(
+ text="Leroy Jenkins",
+ type="reasoning_text",
+ )
+ ],
+ encrypted_content=None,
+ status=None,
+ )
+ formatted_item = construct_chat_message_with_tool_call(item)
+ assert formatted_item["role"] == "assistant"
+ assert formatted_item["reasoning"] == "Leroy Jenkins"
+
+ item = ResponseReasoningItem(
+ id="lol",
+ summary=[
+ Summary(
+ text='Hmm, the user has just started with a simple "Hello,"',
+ type="summary_text",
+ )
+ ],
+ type="reasoning",
+ content=None,
+ encrypted_content=None,
+ status=None,
+ )
+
+ formatted_item = construct_chat_message_with_tool_call(item)
+ assert formatted_item["role"] == "assistant"
+ assert (
+ formatted_item["reasoning"]
+ == 'Hmm, the user has just started with a simple "Hello,"'
+ )
+
+ item = ResponseReasoningItem(
+ id="lol",
+ summary=[],
+ type="reasoning",
+ content=None,
+ encrypted_content="TOP_SECRET_MESSAGE",
+ status=None,
+ )
+ with pytest.raises(ValueError):
+ construct_chat_message_with_tool_call(item)
diff --git a/tests/kernels/attention/test_aiter_flash_attn.py b/tests/kernels/attention/test_aiter_flash_attn.py
index 1dec46e33f22e..8f58c470d217a 100644
--- a/tests/kernels/attention/test_aiter_flash_attn.py
+++ b/tests/kernels/attention/test_aiter_flash_attn.py
@@ -6,6 +6,7 @@ import pytest
import torch
import vllm.v1.attention.backends.rocm_aiter_fa # noqa: F401
+from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
from vllm.platforms import current_platform
NUM_HEADS = [(4, 4), (8, 2)]
@@ -100,6 +101,8 @@ def test_varlen_with_paged_kv(
num_blocks: int,
q_dtype: torch.dtype | None,
) -> None:
+ if not is_flash_attn_varlen_func_available():
+ pytest.skip("flash_attn_varlen_func required to run this test.")
torch.set_default_device("cuda")
current_platform.seed_everything(0)
num_seqs = len(seq_lens)
diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py
index 9662e73321ebe..1a7d5ce0ddc1e 100644
--- a/tests/kernels/attention/test_attention.py
+++ b/tests/kernels/attention/test_attention.py
@@ -13,12 +13,6 @@ from vllm.attention.layer import Attention, MultiHeadAttention
from vllm.platforms import current_platform
from vllm.utils.mem_utils import get_max_shared_memory_bytes
-if not current_platform.is_rocm():
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalCausalMask
-
- from tests.kernels.utils import make_alibi_bias
-
FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
# This will change depending on the compute capability.
# - 512 as a buffer
@@ -448,129 +442,6 @@ def ref_multi_query_kv_attention(
return torch.cat(ref_outputs, dim=0)
-@pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-@pytest.mark.parametrize("head_size", HEAD_SIZES)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@pytest.mark.skipif(
- current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
-)
-@torch.inference_mode()
-def test_multi_query_kv_attention(
- num_seqs: int,
- num_heads: tuple[int, int],
- head_size: int,
- dtype: torch.dtype,
- seed: int,
- device: str,
- use_alibi: bool = False,
-) -> None:
- current_platform.seed_everything(seed)
- torch.set_default_device(device)
- # MAX_SEQ_LEN sometimes causes OOM in the reference implementation.
- # As the xformers library is already tested with its own tests, we can use
- # a smaller MAX_SEQ_LEN here.
- max_len = min(MAX_SEQ_LEN, 4096)
- seq_lens = random.sample(range(1, max_len), num_seqs)
- num_tokens = sum(seq_lens)
-
- scale = float(1.0 / (head_size**0.5))
- num_query_heads, num_kv_heads = num_heads
- qkv = torch.empty(
- num_tokens, num_query_heads + 2 * num_kv_heads, head_size, dtype=dtype
- )
- qkv.uniform_(-scale, scale)
- query, key, value = qkv.split([num_query_heads, num_kv_heads, num_kv_heads], dim=1)
-
- num_queries_per_kv = num_query_heads // num_kv_heads
- if num_queries_per_kv > 1:
- # Handle MQA and GQA
- key = torch.repeat_interleave(key, num_queries_per_kv, dim=1)
- value = torch.repeat_interleave(value, num_queries_per_kv, dim=1)
- alibi_bias = None
- if use_alibi:
- alibi_slopes = torch.randn(num_query_heads, dtype=torch.float)
- attn_bias = make_alibi_bias(alibi_slopes, num_kv_heads, dtype, seq_lens)
- output = torch.empty_like(query)
- start = 0
- # Dynamic sequence length not supported with custom attn_bias.
- for i, seq_len in enumerate(seq_lens):
- end = start + seq_len
- out = xops.memory_efficient_attention_forward(
- query[None, start:end],
- key[None, start:end],
- value[None, start:end],
- attn_bias=attn_bias[i],
- p=0.0,
- scale=scale,
- )
- output[start:end].copy_(out.view_as(query[start:end]))
- start += seq_len
- # xformers.AttentionBias to Tensor for use in reference impl.
- alibi_bias = [
- b.materialize((1, num_query_heads, i, i), device=device).squeeze()
- for b, i in zip(attn_bias, seq_lens)
- ]
- else:
- attn_bias = BlockDiagonalCausalMask.from_seqlens(seq_lens)
- output = xops.memory_efficient_attention_forward(
- query.unsqueeze(0),
- key.unsqueeze(0),
- value.unsqueeze(0),
- attn_bias=attn_bias,
- p=0.0,
- scale=scale,
- )
- output = output.squeeze(0)
-
- cu_seq_lens = [0]
- for seq_len in seq_lens:
- cu_seq_lens.append(cu_seq_lens[-1] + seq_len)
- ref_output = ref_multi_query_kv_attention(
- cu_seq_lens,
- query,
- key,
- value,
- scale,
- alibi_bias,
- dtype,
- )
- atol = get_default_atol(output) if current_platform.is_rocm() else 1e-3
- rtol = get_default_rtol(output) if current_platform.is_rocm() else 1e-5
- torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
-
-
-@pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-@pytest.mark.parametrize("head_size", [64])
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@pytest.mark.skipif(
- current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
-)
-@torch.inference_mode()
-def test_multi_query_kv_attention_with_alibi(
- num_seqs: int,
- num_heads: tuple[int, int],
- head_size: int,
- dtype: torch.dtype,
- seed: int,
- device: str,
-) -> None:
- return test_multi_query_kv_attention(
- num_seqs,
- num_heads,
- head_size,
- dtype,
- seed,
- device,
- use_alibi=True,
- )
-
-
@pytest.mark.parametrize("attention_cls", [Attention, MultiHeadAttention])
def test_num_heads_not_divisble_by_num_kv_heads(attention_cls: type) -> None:
head_size = 64
diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py
index 3b8e939300a27..cd34b520ea71b 100644
--- a/tests/kernels/attention/test_attention_selector.py
+++ b/tests/kernels/attention/test_attention_selector.py
@@ -7,6 +7,7 @@ import pytest
import torch
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
+from vllm.platforms import current_platform
from vllm.platforms.cpu import CpuPlatform
from vllm.platforms.cuda import CudaPlatform
from vllm.platforms.rocm import RocmPlatform
@@ -33,7 +34,7 @@ DEVICE_MLA_BACKENDS = {
}
DEVICE_REGULAR_ATTN_BACKENDS = {
- "cuda": ["XFORMERS", "FLASHINFER", "FLASH_ATTN"],
+ "cuda": ["FLASHINFER", "FLASH_ATTN"],
"hip": ["ROCM_ATTN"],
"cpu": ["CPU_ATTN"],
}
@@ -47,9 +48,11 @@ DEVICE_MLA_BLOCK_SIZES = {
def generate_params():
+ is_rocm = current_platform.is_rocm()
params = []
+ device_list = ["cuda", "cpu"] if not is_rocm else ["hip", "cpu"]
for use_mla in [True, False]:
- for device in ["cuda", "hip", "cpu"]:
+ for device in device_list:
backends = (
DEVICE_MLA_BACKENDS[device]
if use_mla
@@ -204,12 +207,6 @@ def test_env(
)
expected = "FLASHINFER"
assert backend.get_name() == expected
- elif name == "XFORMERS":
- backend = get_attn_backend(
- 32, torch.float16, None, block_size, use_mla=use_mla
- )
- expected = "XFORMERS"
- assert backend.get_name() == expected
elif name == "FLASH_ATTN":
backend = get_attn_backend(
32, torch.float16, None, block_size, use_mla=use_mla
diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py
index f33a27d1fd85a..028e164cb801b 100644
--- a/tests/kernels/attention/test_cache.py
+++ b/tests/kernels/attention/test_cache.py
@@ -68,6 +68,7 @@ def test_copy_blocks(
pytest.skip()
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
# Generate random block mappings where each source block is mapped to two
# destination blocks.
assert 2 * num_mappings <= num_blocks
@@ -152,6 +153,7 @@ def test_reshape_and_cache(
pytest.skip()
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
# Create a random slot mapping.
num_slots = block_size * num_blocks
slot_mapping_lst = random.sample(range(num_slots), num_tokens)
@@ -272,6 +274,7 @@ def test_reshape_and_cache_flash(
) -> None:
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
assert implementation in ["cuda", "triton"]
if implementation == "triton" and kv_cache_layout == "HND":
pytest.skip("Triton implementation only supports NHD layout.")
@@ -593,6 +596,7 @@ def test_concat_and_cache_mla(
) -> None:
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
total_slots = num_blocks * block_size
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
@@ -662,11 +666,14 @@ def test_concat_and_cache_ds_mla(
seed: int,
device: str,
) -> None:
+ if current_platform.is_rocm():
+ pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
if dtype.itemsize != 2:
pytest.skip("ds_mla only supports 16-bit input")
kv_cache_dtype = "fp8_ds_mla"
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
total_slots = num_blocks * block_size
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
@@ -779,6 +786,7 @@ def test_copy_blocks_mla(
) -> None:
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
entry_size = kv_lora_rank + qk_rope_head_dim
@@ -843,6 +851,7 @@ def test_swap_blocks_mla(
) -> None:
current_platform.seed_everything(seed)
torch.set_default_device(device)
+ torch.cuda.set_device(device)
entry_size = kv_lora_rank + qk_rope_head_dim
diff --git a/tests/kernels/attention/test_cascade_flash_attn.py b/tests/kernels/attention/test_cascade_flash_attn.py
index 20f573821b25f..d86041d71febd 100755
--- a/tests/kernels/attention/test_cascade_flash_attn.py
+++ b/tests/kernels/attention/test_cascade_flash_attn.py
@@ -7,11 +7,19 @@ import torch
from vllm.platforms import current_platform
from vllm.v1.attention.backends.flash_attn import cascade_attention, merge_attn_states
-from vllm.vllm_flash_attn import (
- fa_version_unsupported_reason,
- flash_attn_varlen_func,
- is_fa_version_supported,
-)
+
+try:
+ from vllm.vllm_flash_attn import (
+ fa_version_unsupported_reason,
+ flash_attn_varlen_func,
+ is_fa_version_supported,
+ )
+except ImportError:
+ if current_platform.is_rocm():
+ pytest.skip(
+ "vllm_flash_attn is not supported for vLLM on ROCm.",
+ allow_module_level=True,
+ )
NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
HEAD_SIZES = [128, 192, 256]
diff --git a/tests/kernels/attention/test_flash_attn.py b/tests/kernels/attention/test_flash_attn.py
index 6e5468969bf25..bbd5df5419f80 100644
--- a/tests/kernels/attention/test_flash_attn.py
+++ b/tests/kernels/attention/test_flash_attn.py
@@ -6,21 +6,30 @@ import pytest
import torch
from vllm.platforms import current_platform
-from vllm.vllm_flash_attn import (
- fa_version_unsupported_reason,
- flash_attn_varlen_func,
- is_fa_version_supported,
-)
+
+try:
+ from vllm.vllm_flash_attn import (
+ fa_version_unsupported_reason,
+ flash_attn_varlen_func,
+ is_fa_version_supported,
+ )
+except ImportError:
+ if current_platform.is_rocm():
+ pytest.skip(
+ "vllm_flash_attn is not supported for vLLM on ROCm.",
+ allow_module_level=True,
+ )
+
NUM_HEADS = [(4, 4), (8, 2)]
-HEAD_SIZES = [128, 256]
+HEAD_SIZES = [40, 72, 80, 128, 256]
BLOCK_SIZES = [16]
DTYPES = [torch.bfloat16]
QDTYPES = [None, torch.float8_e4m3fn]
# one value large enough to test overflow in index calculation.
# one value small enough to test the schema op check
NUM_BLOCKS = [32768, 2048]
-SOFT_CAPS = [None, 50.0]
+SOFT_CAPS = [None]
SLIDING_WINDOWS = [None, 256]
diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py
index 82ec2ef14e56c..eedeec33e0d45 100644
--- a/tests/kernels/attention/test_flashinfer.py
+++ b/tests/kernels/attention/test_flashinfer.py
@@ -2,12 +2,20 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import flashinfer
import pytest
-import torch
from vllm.platforms import current_platform
+try:
+ import flashinfer
+except ImportError:
+ if current_platform.is_rocm():
+ pytest.skip(
+ "flashinfer is not supported for vLLM on ROCm.", allow_module_level=True
+ )
+
+import torch
+
NUM_HEADS = [(32, 8), (6, 1)]
HEAD_SIZES = [128, 256]
BLOCK_SIZES = [16, 32]
diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py
index 0350136677c6b..d183f67d3919e 100644
--- a/tests/kernels/attention/test_flashinfer_mla_decode.py
+++ b/tests/kernels/attention/test_flashinfer_mla_decode.py
@@ -3,7 +3,6 @@
import pytest
import torch
import torch.nn.functional as F
-from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
from torch import Tensor
from vllm.platforms import current_platform
@@ -15,6 +14,8 @@ if not current_platform.has_device_capability(100):
reason="FlashInfer MLA Requires compute capability of 10 or above.",
allow_module_level=True,
)
+else:
+ from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
def ref_mla(
diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py
index 693b849ebc5d7..98ea40608b468 100644
--- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py
+++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import flashinfer
import pytest
import torch
@@ -16,6 +15,8 @@ if not current_platform.is_device_capability(100):
pytest.skip(
"This TRTLLM kernel requires NVIDIA Blackwell.", allow_module_level=True
)
+else:
+ import flashinfer
FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
FP8_DTYPE = current_platform.fp8_dtype()
diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py
index 9b084f2f660b2..c7662223e1ca5 100644
--- a/tests/kernels/attention/test_merge_attn_states.py
+++ b/tests/kernels/attention/test_merge_attn_states.py
@@ -150,8 +150,8 @@ def test_merge_attn_states(
output_torch = output.clone()
output_lse_torch = output_lse.clone()
total_time_torch_kernel = 0
- start = torch.cuda.Event(enable_timing=True)
- end = torch.cuda.Event(enable_timing=True)
+ start = torch.Event(enable_timing=True)
+ end = torch.Event(enable_timing=True)
# 0. Run the Torch kernel
prefix_lse_torch = prefix_lse.clone()
@@ -188,8 +188,8 @@ def test_merge_attn_states(
output_lse_ref_triton = output_lse.clone()
total_time_triton_kernel = 0
- start = torch.cuda.Event(enable_timing=True)
- end = torch.cuda.Event(enable_timing=True)
+ start = torch.Event(enable_timing=True)
+ end = torch.Event(enable_timing=True)
for _ in range(warmup_times):
merge_attn_states_triton(
diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py
index 183bbf3bf4e03..ae3c63cc62d6b 100644
--- a/tests/kernels/attention/test_mha_attn.py
+++ b/tests/kernels/attention/test_mha_attn.py
@@ -24,10 +24,6 @@ from vllm.platforms.rocm import RocmPlatform
def clear_cache():
"""Clear lru cache to ensure each test case runs without caching."""
_cached_get_attn_backend.cache_clear()
- # Clear xformers availability cache
- import vllm.attention.layer as layer_module
-
- layer_module.USE_XFORMERS_OPS = None
@pytest.mark.parametrize("device", ["cpu", "hip", "cuda"])
@@ -62,38 +58,10 @@ def test_mha_attn_platform(device: str):
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
# Test CUDA with head_size=72 (not divisible by 32)
- # - with upstream FA not available
- # - should use xformers
+ # - should use vLLM's FlashAttention
with (
patch("vllm.attention.layer.current_platform", CudaPlatform()),
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
- patch(
- "vllm.attention.layer.check_upstream_fa_availability",
- return_value=False,
- ),
- ):
- attn = MultiHeadAttention(16, 72, scale=1)
- assert attn.attn_backend == AttentionBackendEnum.XFORMERS
-
- # Test CUDA with head_size=72 (not divisible by 32)
- # - with upstream FA available
- # - should use upstream FA
- with (
- patch("vllm.attention.layer.current_platform", CudaPlatform()),
- patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
- patch(
- "vllm.attention.layer.check_upstream_fa_availability", return_value=True
- ),
- patch.dict(
- "sys.modules",
- {
- "flash_attn": type(
- "MockFlashAttn",
- (),
- {"flash_attn_varlen_func": lambda *args, **kwargs: None},
- )()
- },
- ),
):
attn = MultiHeadAttention(16, 72, scale=1)
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py
index 78cdbbbf7379d..e041e8c8d2ffa 100644
--- a/tests/kernels/attention/test_prefix_prefill.py
+++ b/tests/kernels/attention/test_prefix_prefill.py
@@ -174,11 +174,11 @@ def test_contexted_kv_attention(
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
- b_start_loc = torch.cumsum(torch.tensor([0] + query_lens, dtype=torch.int32), dim=0)
+ b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
max_input_len = MAX_SEQ_LEN
# copy kv to cache
- b_seq_start_loc = torch.cumsum(
- torch.tensor([0] + seq_lens[:-1], dtype=torch.int32), dim=0
+ b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
+ torch.int32
)
for i in range(BS):
for j in range(query_lens[i]):
@@ -417,11 +417,11 @@ def test_contexted_kv_attention_alibi(
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
- b_start_loc = torch.cumsum(torch.tensor([0] + query_lens, dtype=torch.int32), dim=0)
+ b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
max_input_len = MAX_SEQ_LEN
# copy kv to cache
- b_seq_start_loc = torch.cumsum(
- torch.tensor([0] + seq_lens[:-1], dtype=torch.int32), dim=0
+ b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
+ torch.int32
)
for i in range(BS):
for j in range(query_lens[i]):
diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py
index 02b795721f46e..43b242ab2d586 100644
--- a/tests/kernels/core/test_mrope.py
+++ b/tests/kernels/core/test_mrope.py
@@ -5,11 +5,11 @@ from typing import NamedTuple
import pytest
import torch
from packaging.version import Version
-from transformers import AutoConfig
from transformers import __version__ as TRANSFORMERS_VERSION
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.platforms import current_platform
+from vllm.transformers_utils.config import get_config
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -98,8 +98,7 @@ def test_mrope(
atol = model_info.atol
rtol = model_info.rtol
- config = AutoConfig.from_pretrained(model_name)
- config = config.get_text_config()
+ config = get_config(model_name, False).get_text_config()
# get the model config
total_num_kv_heads = config.num_key_value_heads
@@ -113,7 +112,6 @@ def test_mrope(
)
is_neox_style = True
- rope_theta = config.rope_theta
max_position = config.max_position_embeddings
partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
rotary_dim = int(head_dim * partial_rotary_factor)
@@ -122,9 +120,8 @@ def test_mrope(
head_size=head_dim,
rotary_dim=rotary_dim,
max_position=max_position,
- base=rope_theta,
is_neox_style=is_neox_style,
- rope_scaling=config.rope_scaling,
+ rope_parameters=config.rope_parameters,
dtype=dtype,
).to(device=device)
@@ -173,8 +170,7 @@ def test_mrope_torch_compile_tracing(
atol = model_info.atol
rtol = model_info.rtol
- config = AutoConfig.from_pretrained(model_name)
- config = config.get_text_config()
+ config = get_config(model_name, False).get_text_config()
# get the model config
total_num_kv_heads = config.num_key_value_heads
@@ -187,7 +183,6 @@ def test_mrope_torch_compile_tracing(
else config.hidden_size // total_num_heads
)
is_neox_style = True
- rope_theta = config.rope_theta
max_position = config.max_position_embeddings
partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
rotary_dim = int(head_dim * partial_rotary_factor)
@@ -196,9 +191,8 @@ def test_mrope_torch_compile_tracing(
head_size=head_dim,
rotary_dim=rotary_dim,
max_position=max_position,
- base=rope_theta,
is_neox_style=is_neox_style,
- rope_scaling=config.rope_scaling,
+ rope_parameters=config.rope_parameters,
dtype=dtype,
).to(device=device)
diff --git a/tests/kernels/core/test_pos_encoding.py b/tests/kernels/core/test_pos_encoding.py
index c35ee5016ba05..a8ed3825689d3 100644
--- a/tests/kernels/core/test_pos_encoding.py
+++ b/tests/kernels/core/test_pos_encoding.py
@@ -74,7 +74,7 @@ def test_rotary_embedding(
device: str,
use_key: bool,
max_position: int = 8192,
- base: float = 10000,
+ rope_theta: float = 10000,
) -> None:
if rotary_dim is None:
rotary_dim = head_size
@@ -83,7 +83,8 @@ def test_rotary_embedding(
torch.set_default_device(device)
if rotary_dim is None:
rotary_dim = head_size
- rope = get_rope(head_size, rotary_dim, max_position, base, is_neox_style)
+ rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
+ rope = get_rope(head_size, rotary_dim, max_position, is_neox_style, rope_parameters)
rope = rope.to(dtype=dtype, device=torch.get_default_device())
positions = torch.randint(0, max_position, (batch_size, seq_len))
@@ -120,9 +121,9 @@ def test_rotary_embedding(
@torch.inference_mode()
def test_rope_module_cache():
MAX_POSITIONS = [123, 1234]
- BASES = [10000, 1000000]
- ROPE_SCALINGS = (
- None,
+ ROPE_THETAS = [10000, 1000000]
+ ROPE_PARAMETERS = (
+ {"rope_type": "default"},
{"rope_type": "linear", "factor": (1,)},
{"rope_type": "dynamic", "factor": 1},
)
@@ -130,9 +131,9 @@ def test_rope_module_cache():
HEAD_SIZES,
ROTARY_DIMS,
MAX_POSITIONS,
- BASES,
+ ROPE_THETAS,
IS_NEOX_STYLE,
- ROPE_SCALINGS,
+ ROPE_PARAMETERS,
DTYPES,
)
rope_setting_id_map: dict[str, int] = {}
@@ -141,20 +142,20 @@ def test_rope_module_cache():
head_size,
rotary_dim,
max_position,
- base,
- is_neox_stype,
- rope_scaling,
+ rope_theta,
+ is_neox_style,
+ rope_parameters,
dtype,
) = setting
if rotary_dim is None:
rotary_dim = head_size
+ rope_parameters["rope_theta"] = rope_theta
rope = get_rope(
head_size,
rotary_dim,
max_position,
- base,
- is_neox_stype,
- rope_scaling,
+ is_neox_style,
+ rope_parameters,
dtype,
)
# different settings cannot share the same rope module
@@ -168,20 +169,20 @@ def test_rope_module_cache():
head_size,
rotary_dim,
max_position,
- base,
- is_neox_stype,
- rope_scaling,
+ rope_theta,
+ is_neox_style,
+ rope_parameters,
dtype,
) = setting
if rotary_dim is None:
rotary_dim = head_size
+ rope_parameters["rope_theta"] = rope_theta
rope = get_rope(
head_size,
rotary_dim,
max_position,
- base,
- is_neox_stype,
- rope_scaling,
+ is_neox_style,
+ rope_parameters,
dtype,
)
# check if cache take effect
diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py
index 1d925dc1bea8f..d95c22fdf0a5b 100644
--- a/tests/kernels/moe/modular_kernel_tools/common.py
+++ b/tests/kernels/moe/modular_kernel_tools/common.py
@@ -15,7 +15,11 @@ from tests.kernels.quantization.nvfp4_utils import (
)
from tests.kernels.utils import torch_experts
from vllm.config import VllmConfig
-from vllm.distributed import get_dp_group, get_tensor_model_parallel_world_size
+from vllm.distributed import (
+ get_dp_group,
+ get_pcp_group,
+ get_tensor_model_parallel_world_size,
+)
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
@@ -561,6 +565,7 @@ def make_modular_kernel(
# make moe config
moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
tp_size_=get_tensor_model_parallel_world_size(),
+ pcp_size_=get_pcp_group().world_size,
dp_size_=get_dp_group().world_size,
vllm_parallel_config=vllm_config.parallel_config,
)
diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py
index 2285709fa7d60..dab1207d78031 100644
--- a/tests/kernels/moe/test_batched_moe.py
+++ b/tests/kernels/moe/test_batched_moe.py
@@ -39,6 +39,11 @@ MNK_FACTORS = [
NUM_EXPERTS = [8, 64]
TOP_KS = [1, 2, 6]
+DTYPES = [torch.bfloat16]
+
+if not current_platform.is_fp8_fnuz():
+ DTYPES.append(torch.float8_e4m3fn)
+
vllm_config = VllmConfig()
@@ -96,7 +101,7 @@ class BatchedMMTensors:
@pytest.mark.parametrize("max_tokens_per_expert", [32, 224, 512])
@pytest.mark.parametrize("K", [128, 1024])
@pytest.mark.parametrize("N", [128, 1024])
-@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
+@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("block_shape", [None, [128, 128]])
@pytest.mark.parametrize("per_act_token_quant", [False, True])
def test_batched_mm(
@@ -229,7 +234,7 @@ def test_batched_mm(
@pytest.mark.parametrize(("m", "n", "k"), MNK_FACTORS)
@pytest.mark.parametrize("e", NUM_EXPERTS)
@pytest.mark.parametrize("topk", TOP_KS)
-@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
+@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("per_act_token_quant", [False, True])
@pytest.mark.parametrize("block_shape", [None, [128, 128]])
@pytest.mark.parametrize("input_scales", [False])
diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py
index 88db4b3e537c2..b0ff1e64e3219 100644
--- a/tests/kernels/moe/test_block_fp8.py
+++ b/tests/kernels/moe/test_block_fp8.py
@@ -31,6 +31,11 @@ dg_available = has_deep_gemm()
if current_platform.get_device_capability() < (9, 0):
pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True)
+if current_platform.is_fp8_fnuz():
+ pytest.skip(
+ "Tests in this file require float8_e4m3fn and platform does not support",
+ allow_module_level=True,
+ )
vllm_config = VllmConfig()
diff --git a/tests/kernels/moe/test_cutedsl_moe.py b/tests/kernels/moe/test_cutedsl_moe.py
new file mode 100644
index 0000000000000..af1a34d17d48b
--- /dev/null
+++ b/tests/kernels/moe/test_cutedsl_moe.py
@@ -0,0 +1,582 @@
+# 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.has_device_capability(100):
+ pytest.skip(
+ reason="Nvfp4 Requires compute capability of 10 or above.",
+ allow_module_level=True,
+ )
+
+import torch
+from flashinfer import fp4_quantize
+from torch.nn import functional as F
+
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (
+ flashinfer_cutedsl_moe_masked,
+)
+from vllm.utils.flashinfer import (
+ flashinfer_cutedsl_grouped_gemm_nt_masked as cutedsl_gmm_masked,
+)
+from vllm.utils.flashinfer import (
+ scaled_fp4_grouped_quantize,
+)
+
+kE2M1ToFloat = torch.tensor(
+ [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32
+)
+
+FLOAT8_E4M3_MAX = 448.0
+FLOAT4_E2M1_MAX = 6.0
+
+
+def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
+ m_tiles = (m + 128 - 1) // 128
+ f = block_size * 4
+ k_tiles = (k + f - 1) // f
+ tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
+ tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
+ out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
+ return out[0:m, 0:k]
+
+
+def dequantize_nvfp4_to_dtype(
+ tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16
+):
+ """Dequantize the fp4 tensor back to high precision."""
+ # Two fp4 values are packed into one uint8.
+ assert tensor_fp4.dtype == torch.uint8
+ m, packed_k = tensor_fp4.shape
+ k = packed_k * 2
+ tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
+ tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
+ tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
+ tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
+ tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
+
+ # scale the tensor
+ out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
+ return out.to(dtype=dtype)
+
+
+def break_fp4_bytes(a, dtype):
+ assert a.dtype == torch.uint8
+ m, n = a.shape
+
+ # Vectorized nibble processing
+ a_flat = a.flatten()
+ high = (a_flat & 0xF0) >> 4 # Upper nibbles
+ low = a_flat & 0x0F # Lower nibbles
+
+ # Combine nibbles for batch processing
+ combined = torch.stack((low, high), dim=1).flatten()
+
+ # Vectorized sign and magnitude extraction
+ signs = (combined & 0x08).to(torch.bool) # Sign bits
+ abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices
+
+ # Device-aware lookup and sign application
+ kE2M1 = kE2M1ToFloat.to(device=a.device)
+ values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
+
+ # Reshape to final form
+ return values.reshape(m, n * 2).to(dtype=dtype)
+
+
+def generate_balanced_routing(
+ hidden_states: torch.Tensor, num_experts: int, top_k: int
+):
+ """
+ Generate routing weights and topk indices such that every expert is active.
+ Returns routing_weights, topk_idx
+ """
+
+ num_tokens, hidden_dim = hidden_states.shape
+ # num_tokens = batch_size * seq_len
+
+ # First, assign at least one token per expert
+ tokens_per_expert = torch.arange(num_tokens) % num_experts
+ tokens_per_expert = tokens_per_expert[torch.randperm(num_tokens)] # shuffle
+
+ # Each token has top_k experts — start with one guaranteed expert
+ topk_idx = torch.full((num_tokens, top_k), -1, dtype=torch.long)
+ topk_idx[:, 0] = tokens_per_expert
+
+ # For remaining top_k - 1 experts, pick randomly (allowing repeats)
+ if top_k > 1:
+ random_choices = torch.randint(0, num_experts, (num_tokens, top_k - 1))
+ topk_idx[:, 1:] = random_choices
+
+ # Normalize routing weights so each token's weights sum to 1
+ routing_weights = torch.rand(num_tokens, top_k)
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+
+ # Reshape back if needed
+ routing_weights = routing_weights.view(num_tokens, top_k)
+ topk_idx = topk_idx.view(num_tokens, top_k)
+
+ return routing_weights, topk_idx
+
+
+def prepare_inputs(
+ hidden_states: torch.Tensor,
+ router_logits: torch.Tensor,
+ num_experts: int,
+ topk: int,
+):
+ routing_weights, topk_idx = generate_balanced_routing(
+ router_logits, num_experts, topk
+ )
+
+ masked_m = []
+ for i in range(num_experts):
+ mask = topk_idx.view(-1) == i
+ masked_m.append(mask.sum())
+
+ masked_m = torch.tensor(masked_m, dtype=torch.int32)
+ # Intialize the hidden_states_3d with ones instead of empty to avoid nan
+ # issue.
+ hidden_states_3d = torch.ones(
+ (num_experts, max(masked_m), hidden_states.shape[1]), dtype=hidden_states.dtype
+ )
+ for i in range(num_experts):
+ hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i]
+
+ return hidden_states_3d, masked_m, topk_idx, routing_weights
+
+
+MNK_FACTORS = [
+ (2, 1024, 1024),
+ (2, 1024, 1536),
+ (2, 3072, 1024),
+ (2, 3072, 1536),
+ (64, 1024, 1024),
+ (64, 1024, 1536),
+ (64, 3072, 1024),
+ (64, 2048, 1024),
+ (224, 1024, 1024),
+ (224, 1024, 1536),
+]
+
+
+# Reference implementation of torch_moe
+def torch_moe(a, w1, w2, score, topk, expert_map):
+ B, D = a.shape
+ a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+ out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
+ score = torch.softmax(score, dim=-1, dtype=torch.float32)
+ topk_weight, topk_ids = torch.topk(score, topk)
+ topk_weight = topk_weight.view(-1)
+ topk_ids = topk_ids.view(-1)
+ if expert_map is not None:
+ topk_ids = expert_map[topk_ids]
+ for i in range(w1.shape[0]):
+ mask = topk_ids == i
+ if mask.sum():
+ out[mask] = SiluAndMul()(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
+ 0, 1
+ )
+ return (
+ out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
+ ).sum(dim=1)
+
+
+def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids):
+ B, D = a.shape
+ a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+ out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
+
+ topk_weight = topk_weight.view(-1)
+ topk_ids = topk_ids.view(-1)
+
+ for i in range(w1.shape[0]):
+ mask = topk_ids == i
+ if mask.sum():
+ m = w1[i].shape[0]
+ assert m % 2 == 0
+ # Note: w1 and w3 are swapped!
+ w3_expert, w1_expert = w1[i][m // 2 :, :], w1[i][: m // 2, :]
+ inter = F.silu(a[mask] @ w1_expert.t()) * (a[mask] @ w3_expert.t())
+ inter_gs = torch.tensor(1.0).cuda()
+ inter_q, inter_blockscale = fp4_quantize(inter, inter_gs)
+ inter = dequantize_nvfp4_to_dtype(
+ inter_q,
+ inter_blockscale,
+ inter_gs,
+ dtype=inter.dtype,
+ device=inter.device,
+ block_size=16,
+ ).cuda()
+ out[mask] = inter @ w2[i].transpose(0, 1)
+ return (
+ out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
+ ).sum(dim=1)
+
+
+def grouped_gemm_ref(
+ hidden_states_expanded: torch.Tensor,
+ hidden_states_3d: torch.Tensor,
+ weights: torch.Tensor,
+ topk_idx: torch.Tensor,
+ masked_m: torch.Tensor,
+ B: int,
+ topk: int,
+ num_experts: int,
+ *,
+ block_size: int = 16,
+) -> torch.Tensor:
+ """
+ Computes the reference grouped GEMM (fp4 quantized per-expert loop),
+ computes flashinfer grouped GEMM (for scale consistency),
+ and returns ONLY the repacked reference output: out_ref.
+
+ Returns:
+ out_ref: Tensor [num_experts, max_m, n_out]
+ """
+ device_hs = hidden_states_expanded.device
+ device_w = weights.device
+ out_dtype = weights.dtype
+ n_out = weights.shape[1]
+
+ # Flattened reference output (B*topk, n_out)
+ out = torch.zeros((B * topk, n_out), dtype=out_dtype, device=device_w)
+
+ # Per-expert reference compute loop
+ for i in range(num_experts):
+ mask = topk_idx.view(-1) == i
+ if mask.any():
+ lhs = hidden_states_expanded[mask]
+ rhs = weights[i]
+
+ a_amax = lhs.abs().max().to(torch.float32).to(device_hs)
+ b_amax = rhs.abs().max().to(torch.float32).to(device_w)
+
+ a_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax
+ b_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
+
+ lhsq, lhsq_sf = fp4_quantize(lhs, a_gs)
+ rhsq, rhsq_sf = fp4_quantize(rhs, b_gs)
+
+ lhs_in_dtype = dequantize_nvfp4_to_dtype(
+ lhsq,
+ lhsq_sf,
+ a_gs,
+ dtype=lhs.dtype,
+ device=device_hs,
+ block_size=block_size,
+ )
+ rhs_in_dtype = dequantize_nvfp4_to_dtype(
+ rhsq,
+ rhsq_sf,
+ b_gs,
+ dtype=rhs.dtype,
+ device=device_w,
+ block_size=block_size,
+ )
+
+ out[mask] = lhs_in_dtype @ rhs_in_dtype.t()
+
+ # Determine per-expert max_m
+ max_m_val = int(masked_m.max().item())
+
+ # Repack into [num_experts, max_m, n_out]
+ out_ref = torch.zeros(
+ (num_experts, max_m_val, n_out),
+ dtype=out.dtype,
+ device=out.device,
+ )
+ expert_slot = [0] * num_experts
+
+ for i, expert_id in enumerate(topk_idx.view(-1).tolist()):
+ slot = expert_slot[expert_id]
+ if slot < max_m_val:
+ out_ref[expert_id, slot, :] = out[i]
+ expert_slot[expert_id] += 1
+ else:
+ raise IndexError(
+ f"Expert {expert_id} exceeded max slots ({max_m_val}). "
+ "Increase max_m or check masked_m."
+ )
+
+ return out_ref
+
+
+def flashinfer_cutedsl_grouped_gemm_nt_masked(
+ hidden_states: torch.Tensor, # 3d
+ input_global_scale: torch.Tensor, # (l,)
+ weights: torch.Tensor,
+ w_global_scale: torch.Tensor, # (l,)
+ masked_m: torch.Tensor,
+):
+ # hidden_states: [l, m, k]
+ # weights: [l, n, k]
+ aq, aq_sf = scaled_fp4_grouped_quantize(
+ hidden_states,
+ masked_m.to(hidden_states.device),
+ input_global_scale,
+ )
+ num_experts, n, k = weights.shape
+ bq, bq_sf = scaled_fp4_grouped_quantize(
+ weights,
+ torch.full((num_experts,), n, device=weights.device, dtype=torch.int32),
+ w_global_scale,
+ )
+
+ out = torch.zeros(
+ (num_experts, max(masked_m), n), dtype=weights.dtype, device=aq.device
+ )
+ out = out.permute(1, 2, 0) # requirement of kernel
+ sf_vec_size = 16
+ ab_dtype = "float4_e2m1fn"
+ sf_dtype = "float8_e4m3fn"
+ c_dtype = "bfloat16"
+ alpha = 1.0 / (input_global_scale * w_global_scale).to(out.dtype).view(
+ 1, 1, num_experts
+ )
+
+ def get_cute_dtype(input: torch.Tensor) -> str:
+ if input.dtype == torch.bfloat16:
+ return "bfloat16"
+ elif input.dtype == torch.float16:
+ return "float16"
+ elif input.dtype == torch.float32:
+ return "float32"
+ else:
+ raise ValueError(f"Unsupported cute dtype {input.dtype}")
+
+ cutedsl_gmm_masked(
+ (aq, aq_sf),
+ (bq, bq_sf),
+ out,
+ masked_m.to(aq.device),
+ ab_dtype=ab_dtype,
+ sf_dtype=sf_dtype,
+ c_dtype=c_dtype,
+ sf_vec_size=sf_vec_size,
+ alpha=alpha,
+ alpha_dtype=get_cute_dtype(alpha),
+ )
+
+ return out
+
+
+@pytest.mark.parametrize("bs, hidden_dim, inter_dim", [(2, 128, 256), (16, 128, 512)])
+@pytest.mark.parametrize("topk", [1, 2, 4])
+@torch.inference_mode()
+def test_flashinfer_cutedsl_moe_masked(
+ bs: int, hidden_dim: int, inter_dim: int, topk: int
+):
+ torch.manual_seed(42)
+ device = "cuda"
+ num_experts = 8
+ hidden_states = (
+ torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) / 5.0
+ )
+ w1 = (
+ torch.randn(
+ num_experts, 2 * inter_dim, hidden_dim, dtype=torch.bfloat16, device=device
+ )
+ / 10.0
+ )
+ w2 = (
+ torch.randn(
+ num_experts, hidden_dim, inter_dim, dtype=torch.bfloat16, device=device
+ )
+ / 10.0
+ )
+ router_logits = torch.randn(bs, num_experts, dtype=torch.float32)
+
+ hidden_states_expanded = (
+ hidden_states.view(bs, -1, hidden_dim)
+ .repeat(1, topk, 1)
+ .reshape(-1, hidden_dim)
+ )
+ hidden_states_3d, masked_m, topk_idx, routing_weights = prepare_inputs(
+ hidden_states_expanded, router_logits, num_experts, topk
+ )
+
+ w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32).to(w1.device)
+ w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32).to(w2.device)
+ input_global_scale = torch.ones(
+ (num_experts,), dtype=torch.float32, device=hidden_states.device
+ )
+
+ w1_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
+ w2_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
+ a2_global_scale = torch.ones(
+ (num_experts,), dtype=torch.float32, device=hidden_states.device
+ ) # assume intermediate scale is 1.0
+
+ w1_fp4, w1_blockscale = scaled_fp4_grouped_quantize(
+ w1,
+ torch.ones(num_experts, dtype=torch.int32, device=w1.device) * 2 * inter_dim,
+ w1_global_scale,
+ )
+ w2_fp4, w2_blockscale = scaled_fp4_grouped_quantize(
+ w2,
+ torch.ones(num_experts, dtype=torch.int32, device=w2.device) * hidden_dim,
+ w2_global_scale,
+ )
+
+ w1_alpha = 1.0 / (input_global_scale * w1_global_scale)
+ w2_alpha = 1.0 / (a2_global_scale * w2_global_scale)
+
+ out = torch.empty_like(hidden_states_3d)
+ # Note: the 1st dim shouldn't be bs
+ wk = torch.empty(
+ num_experts,
+ hidden_states_3d.shape[1],
+ inter_dim * 2,
+ dtype=hidden_states_3d.dtype,
+ device=hidden_states.device,
+ )
+ flashinfer_cutedsl_moe_masked(
+ hidden_states_3d.to(hidden_states.device),
+ input_global_scale,
+ w1_fp4.permute(2, 0, 1),
+ w1_blockscale,
+ w1_alpha,
+ w2_fp4.permute(2, 0, 1),
+ a2_global_scale,
+ w2_blockscale,
+ w2_alpha,
+ masked_m.to(hidden_states.device),
+ wk,
+ out,
+ )
+
+ # reference
+ a_fp4, a_scale_interleaved = fp4_quantize(hidden_states, input_global_scale)
+ a_in_dtype = dequantize_nvfp4_to_dtype(
+ a_fp4,
+ a_scale_interleaved,
+ input_global_scale,
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
+ block_size=16,
+ )
+ w1_d = torch.empty(
+ (num_experts, 2 * inter_dim, hidden_dim), device=w1.device, dtype=w1.dtype
+ )
+ w2_d = torch.empty(
+ (num_experts, hidden_dim, inter_dim), device=w2.device, dtype=w2.dtype
+ )
+
+ for idx in range(0, num_experts):
+ w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize(
+ w1[idx], w1_global_scale[idx]
+ )
+ w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize(
+ w2[idx], w2_global_scale[idx]
+ )
+ w1_d[idx] = dequantize_nvfp4_to_dtype(
+ w1_fp4_sliced,
+ w1_blockscale_sliced,
+ w1_global_scale[idx],
+ dtype=w1.dtype,
+ device=w1.device,
+ block_size=16,
+ )
+ w2_d[idx] = dequantize_nvfp4_to_dtype(
+ w2_fp4_sliced,
+ w2_blockscale_sliced,
+ w2_global_scale[idx],
+ dtype=w2.dtype,
+ device=w2.device,
+ block_size=16,
+ )
+
+ ref_output = torch_moe_nvfp4(
+ a_in_dtype,
+ w1_d,
+ w2_d,
+ topk,
+ routing_weights.to(a_in_dtype.device),
+ topk_idx.to(a_in_dtype.device),
+ )
+ out_weighted = torch.zeros_like(ref_output, device=out.device, dtype=out.dtype)
+
+ positions = torch.nonzero(masked_m[topk_idx], as_tuple=False)
+ rows, cols = positions[:, 0], positions[:, 1]
+ experts = topk_idx[rows, cols]
+ for i in range(num_experts):
+ mask = experts == i
+ if mask.any():
+ idx = torch.nonzero(mask, as_tuple=False).squeeze(-1)
+ r, c = rows[idx], cols[idx]
+ out_weighted[r] += out[i, : len(r), :] * routing_weights[r, c].to(
+ out.device
+ ).unsqueeze(-1)
+ torch.testing.assert_close(
+ out_weighted.cpu(), ref_output.cpu(), atol=2e-1, rtol=2e-1
+ )
+
+
+@pytest.mark.parametrize(
+ "bs, hidden_dim, inter_dim, topk", [(2, 128, 256, 2), (16, 128, 512, 5)]
+)
+@torch.inference_mode()
+def test_grouped_gemm_nt_masked(
+ bs: int, hidden_dim: int, inter_dim: int, topk: int
+) -> None:
+ torch.manual_seed(42)
+ B = bs
+ D = hidden_dim
+ N = inter_dim
+ # CuteDSL group gemm has issue when not all experts are active.
+ # i.e. masked = [2, 3, 0, 0, 1] where the 2nd and 3rd experts are inactive
+ # see https://github.com/flashinfer-ai/flashinfer/issues/1856
+ num_experts = bs
+ hidden_states = torch.randn(B, D, dtype=torch.bfloat16, device="cuda")
+ weights = torch.randn(num_experts, N, D, dtype=torch.bfloat16, device="cuda")
+ router_logits = torch.randn(B, num_experts, dtype=torch.float32)
+
+ hidden_states_expanded = (
+ hidden_states.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+ )
+ hidden_states_3d, masked_m, topk_idx, _ = prepare_inputs(
+ hidden_states_expanded, router_logits, num_experts, topk
+ )
+
+ a_amax = (
+ hidden_states_3d.abs()
+ .amax(dim=(1, 2))
+ .to(torch.float32)
+ .to(hidden_states.device)
+ )
+ b_amax = weights.abs().amax(dim=(1, 2)).to(torch.float32).to(weights.device)
+ a_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a_amax
+ b_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
+ out_flashinfer = flashinfer_cutedsl_grouped_gemm_nt_masked(
+ hidden_states_3d.to(hidden_states.device), a_gs, weights, b_gs, masked_m
+ )
+ # reference
+ out_ref = grouped_gemm_ref(
+ hidden_states_expanded=hidden_states_expanded,
+ hidden_states_3d=hidden_states_3d,
+ weights=weights,
+ topk_idx=topk_idx,
+ masked_m=masked_m,
+ B=B,
+ topk=topk,
+ num_experts=num_experts,
+ )
+ # Note: just to compare the masked position due to cutedsl may write nan
+ # into unmasked position.
+ for i in range(num_experts):
+ torch.testing.assert_close(
+ out_flashinfer.permute(2, 0, 1)[i, : masked_m[i]],
+ out_ref.to(out_flashinfer.device)[i, : masked_m[i]],
+ atol=1e-1,
+ rtol=1e-1,
+ )
+
+
+if __name__ == "__main__":
+ test_flashinfer_cutedsl_moe_masked(16, 128, 512, 4)
+ test_grouped_gemm_nt_masked(16, 128, 512, 4)
diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py
index 218df4a2632c3..638741e91619b 100644
--- a/tests/kernels/moe/test_flashinfer.py
+++ b/tests/kernels/moe/test_flashinfer.py
@@ -22,7 +22,14 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
from vllm.model_executor.models.llama4 import Llama4MoE
from vllm.platforms import current_platform
-from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+
+try:
+ from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+except ImportError:
+ if current_platform.is_rocm():
+ pytest.skip(
+ "flashinfer not supported for vLLM on ROCm", allow_module_level=True
+ )
if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
90
diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py
index dfd317bcf72f1..98e80ec029777 100644
--- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py
+++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py
@@ -201,7 +201,7 @@ class ModelConfig:
sliding_window: int = 128
initial_context_length: int = 4096
rope_theta: float = 150000.0
- rope_scaling_factor: float = 32.0
+ rope_parameters_factor: float = 32.0
rope_ntk_alpha: float = 1.0
rope_ntk_beta: float = 32.0
@@ -270,6 +270,11 @@ class Case:
@pytest.mark.parametrize("num_token", [2])
@pytest.mark.parametrize("tp", [1, 2, 4, 8])
def test_equiv(num_token, a_dtype, w_dtype, tp):
+ from triton_kernels.tensor_details import layout
+
+ if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"):
+ pytest.skip("make_default_matmul_mxfp4_w_layout not available")
+
M = num_token
E = ModelConfig.num_experts
K = ModelConfig.hidden_size
diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py
index e3b8621b452fa..2a30ef2355529 100644
--- a/tests/kernels/moe/test_modular_kernel_combinations.py
+++ b/tests/kernels/moe/test_modular_kernel_combinations.py
@@ -46,6 +46,12 @@ meets_multi_gpu_requirements = pytest.mark.skipif(
reason="Requires deep_ep or deep_gemm or pplx or flashinfer packages",
)
+if current_platform.is_fp8_fnuz():
+ pytest.skip(
+ "Tests in this file require float8_e4m3fn and platform does not support",
+ allow_module_level=True,
+ )
+
def format_result(verbose, msg, ex=None):
if ex is not None:
diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py
index ba1f657b3ecda..12dd322dccc52 100644
--- a/tests/kernels/moe/test_moe_permute_unpermute.py
+++ b/tests/kernels/moe/test_moe_permute_unpermute.py
@@ -23,6 +23,12 @@ TOP_KS = [2, 6, 8]
EP_SIZE = [1, 4, 16]
current_platform.seed_everything(0)
+if current_platform.is_rocm():
+ pytest.skip(
+ "moe_permute_unpermute_supported is not defined for ROCm",
+ allow_module_level=True,
+ )
+
def torch_permute(
hidden_states: torch.Tensor,
diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
index d6b78dd2c2323..b220205759e2d 100644
--- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
@@ -14,6 +14,12 @@ from vllm.platforms import current_platform
from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT, has_deep_gemm
from vllm.utils.math_utils import cdiv, round_up
+if current_platform.is_fp8_fnuz():
+ pytest.skip(
+ "Tests in this file require float8_e4m3fn and platform does not support",
+ allow_module_level=True,
+ )
+
fp8_dtype = torch.float8_e4m3fn
CASES = [
diff --git a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py
index 7a467e160b784..0ab025dceca40 100644
--- a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+++ b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py
@@ -19,6 +19,12 @@ if current_platform.get_device_capability() < (9, 0):
vllm_config = VllmConfig()
+if current_platform.is_fp8_fnuz():
+ pytest.skip(
+ "Tests in this file require float8_e4m3fn and platform does not support",
+ allow_module_level=True,
+ )
+
def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
"""Matrix multiplication function that supports per-token input
diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py
index e9973c1fcc15e..d0e4f6554a91f 100644
--- a/tests/kernels/quantization/test_block_fp8.py
+++ b/tests/kernels/quantization/test_block_fp8.py
@@ -22,6 +22,7 @@ from vllm.utils.deep_gemm import (
fp8_gemm_nt,
get_col_major_tma_aligned_tensor,
per_block_cast_to_fp8,
+ should_use_deepgemm_for_fp8_linear,
)
from vllm.utils.import_utils import has_deep_gemm
@@ -157,10 +158,6 @@ def test_w8a8_block_fp8_cutlass_matmul():
@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGemm kernels not available.")
@torch.inference_mode()
def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
- # only aligned sizes
- if M % 4 != 0 or K % 128 != 0 or N % 64 != 0:
- pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
-
torch.manual_seed(seed)
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max = fp8_info.max
@@ -168,6 +165,12 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
+ # only aligned sizes are supported by deepgemm
+ if not should_use_deepgemm_for_fp8_linear(
+ output_dtype=out_dtype, weight=B_fp32, supports_deep_gemm=True
+ ):
+ pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
+
A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_size[1])
B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)
diff --git a/tests/kernels/test_cache_kernels.py b/tests/kernels/test_cache_kernels.py
new file mode 100644
index 0000000000000..b5d66b4ede886
--- /dev/null
+++ b/tests/kernels/test_cache_kernels.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for CUDA kernels in cache_kernels.cu."""
+
+import pytest
+import torch
+
+try:
+ from vllm import _custom_ops as ops
+except ImportError:
+ pytest.skip(
+ "Could not import vllm._custom_ops. (pip install -e .)", allow_module_level=True
+ )
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Need CUDA device")
+def test_gather_cache_oob():
+ """
+ Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909).
+ This test constructs a boundary case identified in the issue where
+ seq_starts causes the block_table offset to read out of bounds.
+ """
+
+ batch_size = 1
+ block_size = 64
+ entry_size = 128
+
+ block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
+
+ # This will result in offset = 128 / block_size = 128 / 64 = 2
+ # This will cause the kernel to try to read from
+ # block_table[0, 2], but its size is only 2.
+ seq_starts = torch.tensor([128], dtype=torch.int32, device="cuda")
+
+ seq_len = 65
+ cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
+
+ # src_cache: [num_blocks, block_size, entry_size]
+ num_blocks = 5
+ src_cache = torch.randn(
+ (num_blocks, block_size, entry_size), dtype=torch.float16, device="cuda"
+ )
+
+ dst = torch.empty((seq_len, entry_size), dtype=torch.float16, device="cuda")
+
+ scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
+
+ # Calling the C++ function gather_and_maybe_dequant_cache
+ ops.gather_and_maybe_dequant_cache(
+ src_cache,
+ dst,
+ block_table,
+ cu_seq_lens,
+ batch_size,
+ "auto", # kv_cache_dtype
+ scale,
+ seq_starts,
+ )
+
+ torch.cuda.synchronize()
+ assert True
+
+
+if __name__ == "__main__":
+ pytest.main([__file__])
diff --git a/tests/kernels/utils.py b/tests/kernels/utils.py
index 5d5a26fbfc2cd..9307ef7814a8b 100644
--- a/tests/kernels/utils.py
+++ b/tests/kernels/utils.py
@@ -509,43 +509,6 @@ def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
)
-def make_alibi_bias(
- alibi_slopes: torch.Tensor,
- num_kv_heads: int,
- dtype: torch.dtype,
- seq_lens: list[int],
-) -> list[Any]:
- """Create ALiBi biases compatible with xFormers attention tests."""
- from xformers.ops.fmha.attn_bias import LowerTriangularMaskWithTensorBias
-
- if alibi_slopes is None:
- return [None for _ in seq_lens]
-
- attn_biases: list[Any] = []
- num_heads = alibi_slopes.shape[0]
- assert num_heads >= num_kv_heads, (
- "ALiBi slopes expect at least as many heads as KV heads"
- )
-
- for seq_len in seq_lens:
- bias = torch.arange(seq_len, dtype=dtype, device=alibi_slopes.device)
- bias = bias[None, :] - bias[:, None]
-
- padded_len = (seq_len + 7) // 8 * 8
- bias_tensor = torch.empty(
- 1,
- num_heads,
- seq_len,
- padded_len,
- device=alibi_slopes.device,
- dtype=dtype,
- )[:, :, :, :seq_len].copy_(bias)
- bias_tensor.mul_(alibi_slopes[:, None, None])
- attn_biases.append(LowerTriangularMaskWithTensorBias(bias_tensor))
-
- return attn_biases
-
-
def _make_metadata_tensors(
seq_lens: list[int] | None,
context_lens: list[int] | None,
@@ -649,23 +612,12 @@ def make_kv_cache(
Returns:
- * kv_cache: 2 x num_blocks x (block_size * num_heads * head_size)
- * for backend 'XFORMERS'
* kv_cache: 2 x num_blocks x block_size x num_heads x head_size
* for backend 'FLASH_ATTN'
"""
- if backend == "XFORMERS":
- kv_cache = torch.rand((2, num_blocks, block_size * num_heads * head_size)).to(
- device
- )
- elif backend == "FLASH_ATTN":
- kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(
- device
- )
- else:
- raise ValueError(
- f"Unknown backend value: '{backend}'. Expected 'XFORMERS' or 'FLASH_ATTN'."
- )
+ if backend != "FLASH_ATTN":
+ raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
+ kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(device)
if default_val is not None:
kv_cache[:, :, :] = default_val
return kv_cache
@@ -843,22 +795,14 @@ def assert_actual_matches_ideal(
* output_under_test: actually observed output value
"""
ideal_output = test_params.packed_qkvo.ideal_output
- if backend == "XFORMERS":
- torch.testing.assert_close(
- ideal_output, output_under_test.view_as(ideal_output)
- )
-
- elif backend == "FLASH_ATTN":
- # For FlashAttention override the accuracy thresholds to non default
- # values since we notice a higher difference between the ideal and
- # actual output.
- torch.testing.assert_close(
- ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
- )
- else:
- raise ValueError(
- f"Unknown backend value: '{backend}'. Expected 'XFORMERS' or 'FLASH_ATTN'."
- )
+ if backend != "FLASH_ATTN":
+ raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
+ # For FlashAttention override the accuracy thresholds to non default
+ # values since we notice a higher difference between the ideal and
+ # actual output.
+ torch.testing.assert_close(
+ ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
+ )
# Copied/modified from torch._refs.__init__.py
diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py
index d8ff9339bb49b..9d38ec5422794 100644
--- a/tests/lora/conftest.py
+++ b/tests/lora/conftest.py
@@ -250,6 +250,16 @@ def olmoe_lora_files():
return snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider")
+@pytest.fixture(scope="session")
+def qwen3_lora_files():
+ return snapshot_download(repo_id="charent/self_cognition_Alice")
+
+
+@pytest.fixture(scope="session")
+def llama32_lora_files():
+ return snapshot_download(repo_id="jeeejeee/llama32-3b-text2sql-spider")
+
+
@pytest.fixture
def reset_default_device():
"""
diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py
index dfc45e78e464f..407b29fdd1d58 100644
--- a/tests/lora/test_default_mm_loras.py
+++ b/tests/lora/test_default_mm_loras.py
@@ -5,7 +5,9 @@ Tests for applying default registered multimodal loras.
"""
import os
+import unittest.mock as mock
+import pytest
from huggingface_hub import snapshot_download
from vllm.lora.request import LoRARequest
@@ -114,3 +116,36 @@ def test_default_mm_lora_fails_with_overridden_lora_request(
default_mm_loras={"audio": IMAGE_LORA_PATH},
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
)
+
+
+def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
+ class MockEngineException(Exception):
+ pass
+
+ # Regression test for ensuring default multimodal lora resolution
+ # does not expand the lora req if the prompt type is a string.
+ vllm_runner_kwargs = {
+ **VLLM_RUNNER_BASE_KWARGS,
+ **{"default_mm_loras": {"audio": AUDIO_LORA_PATH}},
+ }
+
+ # Avoid the full generation call since these tests are expensive;
+ # just check what lora request is actually submitted to the engine
+ mock_err = "Engine is mocked for this test"
+
+ with (
+ mock.patch(
+ "vllm.v1.engine.llm_engine.LLMEngine.add_request",
+ side_effect=MockEngineException(mock_err),
+ ) as mock_add_request,
+ vllm_runner(**vllm_runner_kwargs) as vllm_model,
+ ):
+ # Die once we actually submit the request to the engine
+ with pytest.raises(MockEngineException):
+ vllm_model.llm.generate(prompts=AUDIO_PROMPT)
+
+ # Then check to make sure the submitted lora request
+ # and text prompt were zipped together correctly
+ engine_args, engine_kwargs = mock_add_request.call_args
+ assert engine_kwargs["lora_request"] is None
+ assert engine_kwargs["prompt_text"] == AUDIO_PROMPT
diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py
index 91ab4a87c65f8..91c8b861c3c5c 100644
--- a/tests/lora/test_fused_moe_lora_kernel.py
+++ b/tests/lora/test_fused_moe_lora_kernel.py
@@ -1,13 +1,25 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
import random
import pytest
import torch
+from tests.utils import multi_gpu_test
from vllm import _custom_ops as ops
+from vllm.distributed import (
+ init_distributed_environment,
+ initialize_model_parallel,
+ tensor_model_parallel_all_gather,
+ tensor_model_parallel_all_reduce,
+)
+from vllm.distributed.parallel_state import (
+ get_tensor_model_parallel_world_size,
+)
from vllm.lora.ops.triton_ops import fused_moe_lora
from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
@pytest.fixture(autouse=True)
@@ -122,6 +134,8 @@ def use_fused_moe_lora_kernel(
max_loras,
num_experts,
block_size,
+ fully_sharded=False,
+ offset=0,
):
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
@@ -195,10 +209,10 @@ def use_fused_moe_lora_kernel(
config["NUM_STAGES"],
config["SPLIT_K"],
mul_routed_weight,
+ fully_sharded=fully_sharded,
+ offset=offset,
)
- return output
-
def use_torch(
hidden_states,
@@ -317,3 +331,193 @@ def test_fused_moe_lora_kernel(
)
torch.testing.assert_close(output, output2, atol=1e-1, rtol=1e-1)
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("num_tokens", [100])
+@pytest.mark.parametrize("top_k_num", [6])
+@pytest.mark.parametrize("num_experts", [64])
+@pytest.mark.parametrize("max_loras", [4])
+@pytest.mark.parametrize("N", [1408])
+@pytest.mark.parametrize("K", [2048])
+@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("column_parallel", [True, False])
+def test_fused_moe_lora_kernel_fully_sharded(
+ num_tokens,
+ top_k_num,
+ num_experts,
+ max_loras,
+ N,
+ K,
+ max_lora_rank,
+ block_size,
+ dtype,
+ seed,
+ column_parallel,
+):
+ current_platform.seed_everything(seed)
+ # the number of randomly generated sentences.
+ num_sequences = 10
+ # generate data
+ topk_ids, topk_weights, token_lora_mapping = sample_data(
+ num_tokens, num_sequences, max_loras, num_experts, top_k_num
+ )
+
+ def run_torch_spawn(fn, nprocs):
+ torch.multiprocessing.spawn(
+ fn,
+ args=(
+ nprocs,
+ f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
+ dtype,
+ seed,
+ N,
+ K,
+ num_tokens,
+ topk_ids,
+ topk_weights,
+ token_lora_mapping,
+ max_lora_rank,
+ top_k_num,
+ max_loras,
+ num_experts,
+ block_size,
+ column_parallel,
+ ),
+ nprocs=nprocs,
+ )
+
+ run_torch_spawn(use_fused_moe_lora_kernel_tensor_parallel, nprocs=2)
+
+
+def use_fused_moe_lora_kernel_tensor_parallel(
+ local_rank,
+ world_size,
+ init_method,
+ dtype,
+ seed,
+ N,
+ K,
+ num_tokens,
+ topk_ids,
+ topk_weights,
+ token_lora_mapping,
+ max_lora_rank,
+ top_k_num,
+ max_loras,
+ num_experts,
+ block_size,
+ column_parallel,
+):
+ def _get_shard_slice(shard_size):
+ return slice(local_rank * shard_size, (local_rank + 1) * shard_size)
+
+ current_platform.seed_everything(seed)
+
+ device = torch.device(f"cuda:{local_rank}")
+ torch.cuda.set_device(device)
+ torch.set_default_device(device)
+ torch.set_default_dtype(dtype)
+
+ init_distributed_environment(
+ world_size=world_size,
+ rank=local_rank,
+ local_rank=local_rank,
+ distributed_init_method=init_method,
+ )
+ initialize_model_parallel(world_size, 1)
+ tp_size = get_tensor_model_parallel_world_size()
+
+ input_dim = K if column_parallel else N
+ output_dim = N if column_parallel else K
+
+ # init lora weights
+ lora_a = torch.rand(
+ (
+ max_loras,
+ num_experts,
+ max_lora_rank,
+ input_dim,
+ ),
+ dtype=dtype,
+ )
+ lora_b = torch.rand(
+ (
+ max_loras,
+ num_experts,
+ output_dim,
+ max_lora_rank,
+ ),
+ dtype=dtype,
+ )
+
+ hidden_states = torch.rand(
+ (
+ num_tokens,
+ input_dim,
+ ),
+ dtype=dtype,
+ )
+
+ output = torch.zeros((num_tokens, top_k_num, output_dim), dtype=dtype)
+ topk_ids = topk_ids.to(device)
+ topk_weights = topk_weights.to(device)
+ token_lora_mapping = token_lora_mapping.to(device)
+
+ ref_output = use_torch(
+ hidden_states,
+ token_lora_mapping,
+ topk_ids,
+ [lora_a],
+ [lora_b],
+ top_k_num,
+ )
+
+ if column_parallel:
+ # Column parallel (e.g. gate_up_proj): LoRA A is sliced along the rank dim,
+ # and Lora B is sliced along the output dim
+ lora_a_shard_size = max_lora_rank // tp_size
+ lora_a = lora_a[:, :, _get_shard_slice(lora_a_shard_size), :]
+ max_lora_rank = lora_a_shard_size
+ offset = 0
+
+ lora_b_shard_size = output_dim // tp_size
+ lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]
+ output = output[:, :, _get_shard_slice(lora_b_shard_size)].contiguous()
+ else:
+ # Row parallel (e.g. down proj): LoRA A is sliced along the input dim,
+ # and LoRA B is sliced along the output dim
+ lora_a_shard_size = input_dim // tp_size
+ lora_a = lora_a[:, :, :, _get_shard_slice(lora_a_shard_size)]
+ hidden_states = hidden_states[:, _get_shard_slice(lora_a_shard_size)]
+
+ lora_b_shard_size = output_dim // tp_size
+ lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]
+ offset = lora_b_shard_size * local_rank
+
+ use_fused_moe_lora_kernel(
+ topk_ids,
+ topk_weights,
+ token_lora_mapping,
+ max_lora_rank,
+ top_k_num,
+ [lora_a],
+ [lora_b],
+ hidden_states,
+ output,
+ max_loras,
+ num_experts,
+ block_size,
+ fully_sharded=True,
+ offset=offset,
+ )
+
+ if column_parallel:
+ output = tensor_model_parallel_all_gather(output)
+ else:
+ output = tensor_model_parallel_all_reduce(output)
+
+ torch.testing.assert_close(output, ref_output, atol=1e-1, rtol=1e-1)
diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py
index 711d514a39eb3..f4269750feb6b 100644
--- a/tests/lora/test_gptoss_tp.py
+++ b/tests/lora/test_gptoss_tp.py
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+
import vllm
from vllm.lora.request import LoRARequest
@@ -84,14 +86,17 @@ def test_gpt_oss_lora(gptoss20b_lora_files):
@multi_gpu_test(num_gpus=2)
-def test_gpt_oss_lora_tp2(gptoss20b_lora_files):
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
+def test_gpt_oss_lora_tp2(gptoss20b_lora_files, fully_sharded_loras):
llm = vllm.LLM(
MODEL_PATH,
max_model_len=1024,
enable_lora=True,
max_loras=2,
max_lora_rank=8,
+ max_num_seqs=16,
tensor_parallel_size=2,
+ fully_sharded_loras=fully_sharded_loras,
compilation_config=vllm.config.CompilationConfig( # Avoid OOM
cudagraph_specialize_lora=False,
),
diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py
index 8f18f01441932..9df3a07a9e5e9 100644
--- a/tests/lora/test_layers.py
+++ b/tests/lora/test_layers.py
@@ -136,7 +136,6 @@ def populate_loras(
id_to_index: list[int | None],
layer: BaseLayerWithLoRA,
layer_weights: torch.Tensor,
- generate_embeddings_tensor: int = 0,
repeats: int = 1,
) -> tuple[dict[int, LoRALayerWeights], dict[int, list[LoRALayerWeights]]]:
"""This method populates the lora layers with lora weights.
@@ -148,8 +147,6 @@ def populate_loras(
layer: the LoRAlayer to populate.
layer_weights: the PyTorch tensor containing the layer's
weights.
- generate_embeddings_tensor: whether to generate an
- embeddings tensor for each LoRA.
repeats: must only be set for column parallel packed
layers. Indicates the number of loras to compose
together to create a single lora layer.
@@ -171,7 +168,6 @@ def populate_loras(
sublora = DummyLoRAManager(layer_weights.device).init_random_lora(
module_name=f"fake_{i}",
weight=layer_weights,
- generate_embeddings_tensor=generate_embeddings_tensor,
)
sublora.lora_b = sublora.lora_b[
(sublora_len * i) : (sublora_len * (i + 1)), :
@@ -185,7 +181,6 @@ def populate_loras(
slot_idx,
lora_a=lora.lora_a,
lora_b=lora.lora_b,
- embeddings_tensor=lora.embeddings_tensor,
)
lora_dict[lora_id] = lora
@@ -306,7 +301,6 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
id_to_index,
max_loras,
vocab_size,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_embedding(torch.cat(inputs))
@@ -344,7 +338,6 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
id_to_index,
max_loras,
vocab_size,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_embedding(torch.cat(inputs))
@@ -354,149 +347,6 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
-@torch.inference_mode()
-# @pytest.mark.skip(
-# reason="Fails when loras are in any slot other than the first.")
-@pytest.mark.parametrize("num_loras", [1, 2, 4])
-@pytest.mark.parametrize("device", DEVICES)
-@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 128000])
-@pytest.mark.parametrize("stage", STAGES)
-def test_embeddings_with_new_embeddings(
- dist_init, num_loras, device, vocab_size, stage
-) -> None:
- if current_platform.is_cuda_alike():
- torch.cuda.set_device(device)
-
- torch.set_default_device(device)
- max_loras = 8
- punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
- assert check_punica_wrapper(punica_wrapper)
- lora_config = LoRAConfig(
- max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
- )
-
- def create_random_embedding_layer():
- embedding = VocabParallelEmbedding(vocab_size, 256)
- embedding_data = torch.rand_like(embedding.weight.data)
- embedding.weight.data = embedding_data
- embedding.weight.data[vocab_size:, :] = 0
- expanded_embedding = VocabParallelEmbedding(
- vocab_size + lora_config.lora_extra_vocab_size * max_loras,
- 256,
- org_num_embeddings=vocab_size,
- )
- expanded_embedding.weight.data[:vocab_size, :] = embedding_data
- # We need to deepcopy the embedding as it will be modified
- # in place
- lora_embedding = VocabParallelEmbeddingWithLoRA(deepcopy(expanded_embedding))
- lora_embedding.create_lora_weights(max_loras, lora_config)
-
- return expanded_embedding, lora_embedding
-
- for i in range(NUM_RANDOM_SEEDS):
- set_random_seed(i)
-
- id_to_index = get_random_id_to_index(num_loras, max_loras)
- expanded_embedding, lora_embedding = create_random_embedding_layer()
- lora_dict, _ = populate_loras(
- id_to_index,
- layer=lora_embedding,
- layer_weights=torch.zeros(
- (256, vocab_size + lora_config.lora_extra_vocab_size)
- ),
- generate_embeddings_tensor=256,
- )
-
- lora_embedding.set_mapping(punica_wrapper)
- # All embeddings tensors have the same shape.
- embeddings_tensors = [
- lora_dict[id].embeddings_tensor for id in sorted(lora_dict.keys())
- ]
- embeddings_tensor_len = embeddings_tensors[0].shape[0]
-
- # Add empty embeddings_tensors for unoccupied lora slots.
- for _ in range(max_loras - len(embeddings_tensors)):
- embeddings_tensors.append(torch.zeros(embeddings_tensors[0].shape))
-
- inputs, index_mapping, prompt_mapping = create_random_inputs(
- active_lora_ids=list(lora_dict.keys()),
- num_inputs=num_loras * 3,
- input_size=(200,),
- input_range=(1, vocab_size),
- device=device,
- )
- lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
- punica_wrapper.update_metadata(
- lora_mapping,
- id_to_index,
- max_loras,
- vocab_size,
- lora_config.lora_extra_vocab_size,
- )
- original_inputs = deepcopy(inputs)
-
- # Force some of the inputs to be in the extended embeddings range
- # to guarantee that their behavior is tested.
- for input_, original_input_, lora_id in zip(
- inputs, original_inputs, prompt_mapping
- ):
- embedding_id = lora_id - 1
- input_[-1] = vocab_size + (embedding_id * embeddings_tensor_len)
- original_input_[-1] = vocab_size
- input_[-2] = vocab_size + ((embedding_id + 1) * embeddings_tensor_len - 1)
- original_input_[-2] = vocab_size + embeddings_tensor_len - 1
-
- expanded_embedding.weight[
- vocab_size : vocab_size + (embeddings_tensor_len * max_loras)
- ] = torch.cat(embeddings_tensors)
-
- lora_result = lora_embedding(torch.cat(original_inputs))
-
- expected_results: list[torch.Tensor] = []
- for input_, original_input_, lora_id in zip(
- inputs, original_inputs, prompt_mapping
- ):
- lora = lora_dict[lora_id]
- result = expanded_embedding(input_)
- after_a = F.embedding(
- original_input_,
- lora.lora_a.T,
- )
- result += after_a @ lora.lora_b.T
- expected_results.append(result)
- expected_result = torch.cat(expected_results)
-
- rtol, atol = TOLERANCES[lora_result.dtype]
- torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
-
- # Check that resetting the lora weights succeeds
-
- for slot_idx in range(max_loras):
- lora_embedding.reset_lora(slot_idx)
-
- inputs, index_mapping, prompt_mapping = create_random_inputs(
- active_lora_ids=[0],
- num_inputs=num_loras * 3,
- input_size=(200,),
- input_range=(1, vocab_size),
- device=device,
- )
- original_inputs = deepcopy(inputs)
- lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
- punica_wrapper.update_metadata(
- lora_mapping,
- id_to_index,
- max_loras,
- vocab_size,
- lora_config.lora_extra_vocab_size,
- )
- lora_result = lora_embedding(torch.cat(original_inputs))
- expected_result = expanded_embedding(torch.cat(inputs))
-
- rtol, atol = TOLERANCES[lora_result.dtype]
- torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
-
-
@torch.inference_mode()
@pytest.mark.parametrize("num_loras", [1, 2, 4])
@pytest.mark.parametrize("device", DEVICES)
@@ -518,16 +368,13 @@ def test_lm_head_logits_processor(
def _pretest():
linear = ParallelLMHead(
- vocab_size + lora_config.lora_extra_vocab_size,
- 1024,
- vocab_size,
+ num_embeddings=vocab_size,
+ embedding_dim=1024,
params_dtype=torch.float16,
)
linear.weight.data = torch.rand_like(linear.weight.data)
linear.weight.data[:, vocab_size:] = 0
- logits_processor = LogitsProcessor(
- vocab_size + lora_config.lora_extra_vocab_size, vocab_size
- )
+ logits_processor = LogitsProcessor(vocab_size)
lora_logits_processor = LogitsProcessorWithLoRA(
logits_processor, 1024, linear.weight.dtype, linear.weight.device, None
)
@@ -541,15 +388,12 @@ def test_lm_head_logits_processor(
id_to_index = get_random_id_to_index(num_loras, max_loras)
linear, logits_processor, lora_logits_processor = _pretest()
lora_logits_processor.set_mapping(punica_wrapper)
- # NOTE: all the generated loras share the same embeddings tensor.
+
lora_dict, _ = populate_loras(
id_to_index,
layer=lora_logits_processor,
layer_weights=linear.weight,
- generate_embeddings_tensor=1024,
)
- embeddings_tensor = list(lora_dict.values())[0].embeddings_tensor
- embeddings_tensor_len = embeddings_tensor.shape[0]
inputs, index_mapping, prompt_mapping = create_random_inputs(
active_lora_ids=list(lora_dict.keys()),
@@ -565,7 +409,6 @@ def test_lm_head_logits_processor(
id_to_index,
max_loras,
vocab_size,
- lora_config.lora_extra_vocab_size,
)
input_ = torch.rand(20, 1024)
@@ -575,23 +418,16 @@ def test_lm_head_logits_processor(
original_lm_head = deepcopy(linear)
- linear.weight[
- logits_processor.org_vocab_size : logits_processor.org_vocab_size
- + embeddings_tensor_len
- ] = embeddings_tensor
-
- logits_processor.org_vocab_size = vocab_size + lora_config.lora_extra_vocab_size
expected_results: list[torch.Tensor] = []
for input_, lora_id in zip(inputs, prompt_mapping):
lora = lora_dict[lora_id]
result = logits_processor._get_logits(
hidden_states=input_, lm_head=linear, embedding_bias=None
)
- result[:, vocab_size + embeddings_tensor_len :] = float("-inf")
+
result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling
expected_results.append(result)
expected_result = torch.cat(expected_results)
- logits_processor.org_vocab_size = vocab_size
# Check that resetting the lora weights succeeds
@@ -612,7 +448,6 @@ def test_lm_head_logits_processor(
id_to_index,
max_loras,
vocab_size,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_logits_processor._get_logits(
@@ -694,7 +529,6 @@ def test_linear_replicated(
id_to_index,
max_loras,
512,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_linear(torch.cat(inputs))[0]
@@ -726,7 +560,10 @@ def test_linear_replicated(
lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
punica_wrapper.update_metadata(
- lora_mapping, id_to_index, max_loras, 512, lora_config.lora_extra_vocab_size
+ lora_mapping,
+ id_to_index,
+ max_loras,
+ 512,
)
lora_result = lora_linear(torch.cat(inputs))[0]
@@ -817,7 +654,6 @@ def test_linear_parallel(
id_to_index,
max_loras,
512,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_linear(torch.cat(inputs))[0]
@@ -849,7 +685,10 @@ def test_linear_parallel(
lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
punica_wrapper.update_metadata(
- lora_mapping, id_to_index, max_loras, 512, lora_config.lora_extra_vocab_size
+ lora_mapping,
+ id_to_index,
+ max_loras,
+ 512,
)
lora_result = lora_linear(torch.cat(inputs))[0]
@@ -963,7 +802,6 @@ def test_column_parallel_packed(
id_to_index,
max_loras,
512,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_linear(torch.cat(inputs))[0]
@@ -1000,7 +838,6 @@ def test_column_parallel_packed(
id_to_index,
max_loras,
512,
- lora_config.lora_extra_vocab_size,
)
lora_result = lora_linear(torch.cat(inputs))[0]
diff --git a/tests/lora/test_llama_tp.py b/tests/lora/test_llama_tp.py
index 7bbd1e364d19e..18704fa6e45de 100644
--- a/tests/lora/test_llama_tp.py
+++ b/tests/lora/test_llama_tp.py
@@ -13,17 +13,27 @@ from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
-MODEL_PATH = "meta-llama/Llama-2-7b-hf"
+PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|>
+I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
+"
+##Instruction:
+candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
+Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
+The People_ID of candidate is the foreign key of People_ID of people.
+###Input:
+{context}
+###Response:<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+""" # noqa: E501
EXPECTED_LORA_OUTPUT = [
- " SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport' ", # noqa: E501
- " SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone' ",
- " SELECT one_mora FROM table_name_95 WHERE gloss = 'low tone mora with a gloss of /˩okiru/' [òkìɽɯ́] AND accented_mora = 'low tone mora with a gloss of /˩okiru/' [òkìɽɯ́] ", # noqa: E501
- " SELECT sex FROM people WHERE people_id IN (SELECT people_id FROM candidate GROUP BY sex ORDER BY COUNT(people_id) DESC LIMIT 1) ", # noqa: E501
- " SELECT pick FROM table_name_60 WHERE former_wnba_team = 'Minnesota Lynx' ",
- " SELECT womens_doubles FROM table_28138035_4 WHERE mens_singles = 'Werner Schlager' ", # noqa: E501
+ "SELECT count(*) FROM candidate",
+ "SELECT count(*) FROM candidate",
+ "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
+ "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
]
+MODEL_PATH = "meta-llama/Llama-3.2-3B-Instruct"
+
def do_sample(
llm: vllm.LLM,
@@ -32,18 +42,19 @@ def do_sample(
tensorizer_config_dict: dict | None = None,
) -> list[str]:
prompts = [
- "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501
- "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]", # noqa: E501
- "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_95 (one_mora VARCHAR, gloss VARCHAR, accented_mora VARCHAR)\n\n question: What is the one mora for a low tone mora with a gloss of /˩okiru/ [òkìɽɯ́]? [/user] [assistant]", # noqa: E501
- "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE candidate (people_id VARCHAR, unsure_rate INTEGER); CREATE TABLE people (sex VARCHAR, people_id VARCHAR)\n\n question: which gender got the highest average uncertain ratio. [/user] [assistant]", # noqa: E501
- "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_60 (pick INTEGER, former_wnba_team VARCHAR)\n\n question: What pick was a player that previously played for the Minnesota Lynx? [/user] [assistant]", # noqa: E501
- "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_28138035_4 (womens_doubles VARCHAR, mens_singles VARCHAR)\n\n question: Name the women's doubles for werner schlager [/user] [assistant]", # noqa: E501
+ PROMPT_TEMPLATE.format(context="How many candidates are there?"),
+ PROMPT_TEMPLATE.format(context="Count the number of candidates."),
+ PROMPT_TEMPLATE.format(
+ context="Which poll resource provided the most number of candidate information?" # noqa: E501
+ ),
+ PROMPT_TEMPLATE.format(
+ context="Return the poll resource associated with the most candidates."
+ ),
]
sampling_params = vllm.SamplingParams(
- temperature=0, max_tokens=256, skip_special_tokens=False, stop=["[/assistant]"]
+ temperature=0, max_tokens=64, stop=["<|im_end|>"]
)
-
if tensorizer_config_dict is not None:
outputs = llm.generate(
prompts,
@@ -75,13 +86,15 @@ def do_sample(
return generated_texts
-def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None = None):
+def generate_and_test(
+ llm, llama32_lora_files, tensorizer_config_dict: dict | None = None
+):
print("lora adapter created")
print("lora 1")
assert (
do_sample(
llm,
- sql_lora_files,
+ llama32_lora_files,
tensorizer_config_dict=tensorizer_config_dict,
lora_id=1,
)
@@ -92,7 +105,7 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =
assert (
do_sample(
llm,
- sql_lora_files,
+ llama32_lora_files,
tensorizer_config_dict=tensorizer_config_dict,
lora_id=2,
)
@@ -104,51 +117,52 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =
@create_new_process_for_each_test()
@pytest.mark.parametrize("cudagraph_specialize_lora", [True, False])
-def test_llama_lora(sql_lora_files, cudagraph_specialize_lora: bool):
+def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool):
llm = vllm.LLM(
MODEL_PATH,
- tokenizer=sql_lora_files,
enable_lora=True,
# also test odd max_num_seqs
- max_num_seqs=13,
+ max_num_seqs=7,
+ max_model_len=1024,
max_loras=4,
compilation_config=vllm.config.CompilationConfig(
cudagraph_specialize_lora=cudagraph_specialize_lora,
),
)
- generate_and_test(llm, sql_lora_files)
+ generate_and_test(llm, llama32_lora_files)
@multi_gpu_test(num_gpus=4)
-def test_llama_lora_tp4(sql_lora_files):
+def test_llama_lora_tp4(llama32_lora_files):
llm = vllm.LLM(
MODEL_PATH,
- tokenizer=sql_lora_files,
enable_lora=True,
- max_num_seqs=16,
+ max_num_seqs=7,
+ max_model_len=1024,
max_loras=4,
tensor_parallel_size=4,
)
- generate_and_test(llm, sql_lora_files)
+ generate_and_test(llm, llama32_lora_files)
@multi_gpu_test(num_gpus=4)
-def test_llama_lora_tp4_fully_sharded_loras(sql_lora_files):
+def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files):
llm = vllm.LLM(
MODEL_PATH,
- tokenizer=sql_lora_files,
enable_lora=True,
- max_num_seqs=16,
+ max_num_seqs=8,
max_loras=4,
+ max_model_len=1024,
tensor_parallel_size=4,
fully_sharded_loras=True,
)
- generate_and_test(llm, sql_lora_files)
+ generate_and_test(llm, llama32_lora_files)
@multi_gpu_test(num_gpus=2)
def test_tp2_serialize_and_deserialize_lora(
- tmp_path, sql_lora_files, sql_lora_huggingface_id
+ tmp_path,
+ llama32_lora_files,
):
# Run the tensorizing of the LoRA adapter and the model in a subprocess
# to guarantee cleanup
@@ -157,7 +171,7 @@ def test_tp2_serialize_and_deserialize_lora(
model_name = "model-rank-%03d.tensors"
model_ref = MODEL_PATH
- lora_path = sql_lora_huggingface_id
+ lora_path = llama32_lora_files
suffix = "test"
try:
result = subprocess.run(
@@ -195,12 +209,12 @@ def test_tp2_serialize_and_deserialize_lora(
loaded_llm = LLM(
model=model_ref,
- tokenizer=sql_lora_files,
load_format="tensorizer",
enable_lora=True,
enforce_eager=True,
model_loader_extra_config=tensorizer_config,
- max_num_seqs=13,
+ max_num_seqs=7,
+ max_model_len=1024,
tensor_parallel_size=2,
max_loras=2,
)
@@ -211,7 +225,7 @@ def test_tp2_serialize_and_deserialize_lora(
print("lora 1")
assert (
do_sample(
- loaded_llm, sql_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
+ loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
)
== EXPECTED_LORA_OUTPUT
)
diff --git a/tests/lora/test_lora_functions.py b/tests/lora/test_lora_functions.py
index e914393fee8aa..1c692630284d0 100644
--- a/tests/lora/test_lora_functions.py
+++ b/tests/lora/test_lora_functions.py
@@ -13,8 +13,8 @@ from vllm.entrypoints.openai.api_server import (
from vllm.lora.request import LoRARequest
from vllm.v1.engine.llm_engine import LLMEngine
-MODEL_PATH = "meta-llama/Llama-2-7b-hf"
-LORA_MODULE_PATH = "yard1/llama-2-7b-sql-lora-test"
+MODEL_PATH = "Qwen/Qwen3-0.6B"
+LORA_MODULE_PATH = "charent/self_cognition_Alice"
LORA_RANK = 8
diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py
index e7816031142e3..24d4dfca46d62 100644
--- a/tests/lora/test_lora_manager.py
+++ b/tests/lora/test_lora_manager.py
@@ -48,9 +48,6 @@ DEFAULT_DTYPE = torch.get_default_dtype()
@pytest.mark.parametrize("device", DEVICES)
def test_from_lora_tensors(sql_lora_files, device):
tensors = load_file(os.path.join(sql_lora_files, "adapter_model.safetensors"))
- new_embeddings = load_file(
- os.path.join(sql_lora_files, "new_embeddings.safetensors")
- )
peft_helper = PEFTHelper.from_local_dir(
sql_lora_files, max_position_embeddings=4096
@@ -60,7 +57,6 @@ def test_from_lora_tensors(sql_lora_files, device):
tensors,
peft_helper=peft_helper,
device=device,
- embeddings=new_embeddings,
embedding_modules=EMBEDDING_MODULES,
embedding_padding_modules=EMBEDDING_PADDING_MODULES,
)
@@ -76,18 +72,6 @@ def test_from_lora_tensors(sql_lora_files, device):
f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"
)
assert lora.lora_a.shape[0] == 8
- embeddings_module = next(
- (k for k in EMBEDDING_MODULES if k in module_name), None
- )
- if embeddings_module:
- assert torch.equal(
- lora.embeddings_tensor,
- new_embeddings[EMBEDDING_MODULES[embeddings_module]].to(
- device=lora.embeddings_tensor.device
- ),
- )
- else:
- assert lora.embeddings_tensor is None
def create_lora(
@@ -552,9 +536,7 @@ def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path
worker_adapter_manager = WorkerLoRAManager(
vllm_config, device, EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES
)
- worker_adapter_manager.vocab_size = (
- dummy_model_gate_up.unpadded_vocab_size - lora_config.lora_extra_vocab_size
- )
+ worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
worker_adapter_manager.create_lora_manager(dummy_model_gate_up)
dummy_lora_files = f"{tmp_path}/lora_adapter"
diff --git a/tests/lora/test_minicpmv_tp.py b/tests/lora/test_minicpmv_tp.py
index 1cf8ed602b6a4..e430826461a14 100644
--- a/tests/lora/test_minicpmv_tp.py
+++ b/tests/lora/test_minicpmv_tp.py
@@ -57,10 +57,6 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
return generated_texts
-@pytest.mark.xfail(
- current_platform.is_rocm(),
- reason="MiniCPM-V dependency xformers incompatible with ROCm",
-)
def test_minicpmv_lora(minicpmv_lora_files):
llm = vllm.LLM(
MODEL_PATH,
@@ -84,10 +80,6 @@ def test_minicpmv_lora(minicpmv_lora_files):
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
-@pytest.mark.xfail(
- current_platform.is_rocm(),
- reason="MiniCPM-V dependency xformers incompatible with ROCm",
-)
@multi_gpu_test(num_gpus=4)
def test_minicpmv_tp4_wo_fully_sharded_loras(minicpmv_lora_files):
llm = vllm.LLM(
@@ -108,10 +100,6 @@ def test_minicpmv_tp4_wo_fully_sharded_loras(minicpmv_lora_files):
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
-@pytest.mark.xfail(
- current_platform.is_rocm(),
- reason="MiniCPM-V dependency xformers incompatible with ROCm",
-)
@multi_gpu_test(num_gpus=4)
def test_minicpmv_tp4_fully_sharded_loras(minicpmv_lora_files):
llm = vllm.LLM(
diff --git a/tests/lora/test_olmoe_tp.py b/tests/lora/test_olmoe_tp.py
index e659c1e1a9a07..e3c9816625ba7 100644
--- a/tests/lora/test_olmoe_tp.py
+++ b/tests/lora/test_olmoe_tp.py
@@ -2,6 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+
import vllm
from vllm.lora.request import LoRARequest
@@ -111,8 +113,9 @@ def test_olmoe_lora_mixed(olmoe_lora_files):
generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
@multi_gpu_test(num_gpus=2)
-def test_olmoe_lora_tp2(olmoe_lora_files):
+def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
llm = vllm.LLM(
MODEL_PATH,
max_model_len=1024,
@@ -122,14 +125,16 @@ def test_olmoe_lora_tp2(olmoe_lora_files):
trust_remote_code=True,
enable_chunked_prefill=True,
tensor_parallel_size=2,
+ fully_sharded_loras=fully_sharded_loras,
)
generate_and_test(llm, olmoe_lora_files, lora_id=1)
generate_and_test(llm, olmoe_lora_files, lora_id=2)
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
@multi_gpu_test(num_gpus=4)
-def test_olmoe_lora_tp4(olmoe_lora_files):
+def test_olmoe_lora_tp4(olmoe_lora_files, fully_sharded_loras):
llm = vllm.LLM(
MODEL_PATH,
max_model_len=1024,
@@ -139,6 +144,7 @@ def test_olmoe_lora_tp4(olmoe_lora_files):
trust_remote_code=True,
enable_chunked_prefill=True,
tensor_parallel_size=4,
+ fully_sharded_loras=fully_sharded_loras,
)
generate_and_test(llm, olmoe_lora_files, lora_id=1)
diff --git a/tests/lora/test_qwen2vl.py b/tests/lora/test_qwen2vl.py
index 1800ca107a426..7d8c940100ca4 100644
--- a/tests/lora/test_qwen2vl.py
+++ b/tests/lora/test_qwen2vl.py
@@ -2,12 +2,9 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
-import pytest
-
import vllm
from vllm.assets.image import ImageAsset
from vllm.lora.request import LoRARequest
-from vllm.platforms import current_platform
from vllm.sampling_params import BeamSearchParams
@@ -142,10 +139,6 @@ QWEN2VL_MODEL_PATH = "Qwen/Qwen2-VL-2B-Instruct"
QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"
-@pytest.mark.xfail(
- current_platform.is_rocm(),
- reason="Qwen2-VL dependency xformers incompatible with ROCm",
-)
def test_qwen2vl_lora(qwen2vl_lora_files):
"""Test Qwen 2.0 VL model with LoRA"""
config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
@@ -156,10 +149,6 @@ def test_qwen2vl_lora(qwen2vl_lora_files):
tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
-@pytest.mark.xfail(
- current_platform.is_rocm(),
- reason="Qwen2-VL dependency xformers incompatible with ROCm",
-)
def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
"""Test Qwen 2.0 VL model with LoRA through beam search."""
config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
@@ -178,10 +167,6 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
)
-@pytest.mark.xfail(
- current_platform.is_rocm(),
- reason="Qwen2.5-VL dependency xformers incompatible with ROCm",
-)
def test_qwen25vl_lora(qwen25vl_lora_files):
"""Test Qwen 2.5 VL model with LoRA"""
config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
diff --git a/tests/lora/test_worker.py b/tests/lora/test_worker.py
index c97f8debd1b9a..b163559a9414d 100644
--- a/tests/lora/test_worker.py
+++ b/tests/lora/test_worker.py
@@ -20,11 +20,12 @@ from vllm.lora.models import LoRAMapping
from vllm.lora.request import LoRARequest
from vllm.v1.worker.gpu_worker import Worker
+MODEL_PATH = "Qwen/Qwen3-0.6B"
NUM_LORAS = 16
@patch.dict(os.environ, {"RANK": "0"})
-def test_worker_apply_lora(sql_lora_files):
+def test_worker_apply_lora(qwen3_lora_files):
def set_active_loras(worker: Worker, lora_requests: list[LoRARequest]):
lora_mapping = LoRAMapping([], [])
@@ -34,9 +35,10 @@ def test_worker_apply_lora(sql_lora_files):
vllm_config = VllmConfig(
model_config=ModelConfig(
- "meta-llama/Llama-2-7b-hf",
+ MODEL_PATH,
seed=0,
dtype="float16",
+ max_model_len=127,
enforce_eager=True,
),
load_config=LoadConfig(
@@ -73,7 +75,7 @@ def test_worker_apply_lora(sql_lora_files):
assert worker.list_loras() == set()
lora_requests = [
- LoRARequest(str(i + 1), i + 1, sql_lora_files) for i in range(NUM_LORAS)
+ LoRARequest(str(i + 1), i + 1, qwen3_lora_files) for i in range(NUM_LORAS)
]
set_active_loras(worker, lora_requests)
diff --git a/tests/lora/utils.py b/tests/lora/utils.py
index d30b77f094665..6aba5299b5829 100644
--- a/tests/lora/utils.py
+++ b/tests/lora/utils.py
@@ -28,7 +28,6 @@ class DummyLoRAManager:
module_name: str,
weight: torch.Tensor,
rank: int = 8,
- generate_embeddings_tensor: int = 0,
):
lora = LoRALayerWeights(
module_name,
@@ -41,13 +40,6 @@ class DummyLoRAManager:
[weight.shape[0], rank], dtype=weight.dtype, device=self._device
),
)
- if generate_embeddings_tensor:
- lora.embeddings_tensor = torch.rand(
- 5,
- generate_embeddings_tensor,
- dtype=weight.dtype,
- device=self._device,
- )
self.set_module_lora(module_name, lora)
return lora
diff --git a/tests/model_executor/model_loader/fastsafetensors_loader/test_fastsafetensors_loader.py b/tests/model_executor/model_loader/fastsafetensors_loader/test_fastsafetensors_loader.py
index f154df6dfc232..c5b3c731ffc64 100644
--- a/tests/model_executor/model_loader/fastsafetensors_loader/test_fastsafetensors_loader.py
+++ b/tests/model_executor/model_loader/fastsafetensors_loader/test_fastsafetensors_loader.py
@@ -19,7 +19,8 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
@pytest.mark.skipif(
- not current_platform.is_cuda(), reason="fastsafetensors requires CUDA/NVIDIA GPUs"
+ not current_platform.is_cuda_alike(),
+ reason="fastsafetensors requires NVIDIA/AMD GPUs",
)
def test_model_loader_download_files(vllm_runner):
with vllm_runner(test_model, load_format="fastsafetensors") as llm:
diff --git a/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py b/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py
index bd216f0e41a47..1975eb61b25da 100644
--- a/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py
+++ b/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py
@@ -17,7 +17,8 @@ from vllm.platforms import current_platform
@pytest.mark.skipif(
- not current_platform.is_cuda(), reason="fastsafetensors requires CUDA/NVIDIA GPUs"
+ not current_platform.is_cuda_alike(),
+ reason="fastsafetensors requires NVIDIA/AMD GPUs",
)
def test_fastsafetensors_model_loader():
with tempfile.TemporaryDirectory() as tmpdir:
diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py
index 5bb841bf2fa0e..cf06b000efb51 100644
--- a/tests/model_executor/model_loader/test_sharded_state_loader.py
+++ b/tests/model_executor/model_loader/test_sharded_state_loader.py
@@ -60,18 +60,9 @@ def llama_3p2_1b_files():
def _run_writer(input_dir, output_dir, weights_patterns, **kwargs):
llm_sharded_writer = LLM(model=input_dir, **kwargs)
- # Check which engine version is being used
- is_v1_engine = hasattr(llm_sharded_writer.llm_engine, "engine_core")
+
# Dump worker states to output directory
- if is_v1_engine:
- # For V1 engine, we need to use engine_core.save_sharded_state
- print("Using V1 engine save path")
- llm_sharded_writer.llm_engine.engine_core.save_sharded_state(path=output_dir)
- else:
- # For V0 engine
- print("Using V0 engine save path")
- model_executor = llm_sharded_writer.llm_engine.model_executor
- model_executor.save_sharded_state(path=output_dir)
+ llm_sharded_writer.llm_engine.engine_core.save_sharded_state(path=output_dir)
# Copy metadata files to output directory
for file in os.listdir(input_dir):
diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py
index 0cdb7c9a603f2..df6c2cab7814b 100644
--- a/tests/models/language/generation/test_common.py
+++ b/tests/models/language/generation/test_common.py
@@ -10,13 +10,6 @@ from ....utils import large_gpu_mark
from ...registry import HF_EXAMPLE_MODELS
from ...utils import check_logprobs_close
-# These have unsupported head_dim for FA. We do not
-# have a clean way to fall back, so we fail with
-# a clear msg when it happens.
-# https://github.com/vllm-project/vllm/issues/14524
-# NOTE(woosuk): Skipping these tests until V1 supports them.
-# REQUIRES_V0 = ["microsoft/phi-2", "stabilityai/stablelm-3b-4e1t"]
-
# This list contains the model that are using AITER kernel.
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be
diff --git a/tests/models/language/generation/test_mistral.py b/tests/models/language/generation/test_mistral.py
index 0ae83ec16020a..80e337d570a36 100644
--- a/tests/models/language/generation/test_mistral.py
+++ b/tests/models/language/generation/test_mistral.py
@@ -208,7 +208,7 @@ def test_mistral_format(
with vllm_runner(
model,
dtype=dtype,
- tokenizer_mode="auto",
+ tokenizer_mode="hf",
load_format="safetensors",
config_format="hf",
) as hf_format_model:
diff --git a/tests/models/language/pooling/test_mm_classifier_conversion.py b/tests/models/language/pooling/test_mm_classifier_conversion.py
index 2482452645ef1..a31a771238e26 100644
--- a/tests/models/language/pooling/test_mm_classifier_conversion.py
+++ b/tests/models/language/pooling/test_mm_classifier_conversion.py
@@ -75,7 +75,7 @@ def test_gemma_multimodal(
{
"type": "image_url",
"image_url": {
- "url": "https://upload.wikimedia.org/wikipedia/commons/c/c6/Set_of_fourteen_side_chairs_MET_DP110780.jpg"
+ "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/red_chair.jpg"
},
},
{"type": "text", "text": "A fine 19th century piece of furniture."},
diff --git a/tests/models/language/pooling/test_nomic_max_model_len.py b/tests/models/language/pooling/test_nomic_max_model_len.py
index 88f088c603276..d6216a87a229e 100644
--- a/tests/models/language/pooling/test_nomic_max_model_len.py
+++ b/tests/models/language/pooling/test_nomic_max_model_len.py
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# ruff: noqa: SIM117
+from typing import Any
+
import pytest
from ...utils import EmbedModelInfo
@@ -79,8 +81,8 @@ def test_set_max_model_len_illegal(model_info, vllm_runner):
@pytest.mark.parametrize("model_info", MODELS)
def test_use_rope_scaling_legal(model_info, vllm_runner):
hf_overrides = {
- "rope_theta": rope_theta,
- "rope_scaling": {
+ "rope_parameters": {
+ "rope_theta": rope_theta,
"rope_type": "yarn",
"factor": factor,
"original_max_position_embeddings": original_max_position_embeddings,
@@ -96,9 +98,9 @@ def test_use_rope_scaling_legal(model_info, vllm_runner):
@pytest.mark.parametrize("model_info", MODELS)
def test_use_rope_scaling_illegal(model_info, vllm_runner):
- hf_overrides = {
- "rope_theta": rope_theta,
- "rope_scaling": {
+ hf_overrides: dict[str, Any] = {
+ "rope_parameters": {
+ "rope_theta": rope_theta,
"rope_type": "yarn",
"factor": factor,
"original_max_position_embeddings": original_max_position_embeddings,
@@ -115,8 +117,8 @@ def test_use_rope_scaling_illegal(model_info, vllm_runner):
pass
hf_overrides = {
- "rope_theta": rope_theta,
- "rope_scaling": {
+ "rope_parameters": {
+ "rope_theta": rope_theta,
"rope_type": "yarn",
"factor": factor,
"original_max_position_embeddings": original_max_position_embeddings,
diff --git a/tests/models/language/pooling_mteb_test/mteb_utils.py b/tests/models/language/pooling_mteb_test/mteb_utils.py
index 0384ff82790f0..189cdbae99dcd 100644
--- a/tests/models/language/pooling_mteb_test/mteb_utils.py
+++ b/tests/models/language/pooling_mteb_test/mteb_utils.py
@@ -2,12 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import tempfile
-from collections.abc import Sequence
import mteb
import numpy as np
import requests
import torch
+from mteb.models import ModelMeta
+from mteb.types import Array
+from torch.utils.data import DataLoader
import tests.ci_envs as ci_envs
from tests.models.utils import (
@@ -27,24 +29,47 @@ MTEB_EMBED_TOL = 1e-4
# See #19344
MTEB_RERANK_TASKS = ["NFCorpus"]
-MTEB_RERANK_LANGS = ["en"]
+MTEB_RERANK_LANGS = ["eng"]
MTEB_RERANK_TOL = 2e-3
+_empty_model_meta = ModelMeta(
+ loader=None,
+ name="vllm/model",
+ revision="1",
+ release_date=None,
+ languages=None,
+ framework=[],
+ similarity_fn_name=None,
+ n_parameters=None,
+ memory_usage_mb=None,
+ max_tokens=None,
+ embed_dim=None,
+ license=None,
+ open_weights=None,
+ public_training_code=None,
+ public_training_data=None,
+ use_instructions=None,
+ training_datasets=None,
+ modalities=["text"], # 'image' can be added to evaluate multimodal models
+)
+
+
+class VllmMtebEncoder(mteb.EncoderProtocol):
+ mteb_model_meta = _empty_model_meta
-class VllmMtebEncoder(mteb.Encoder):
def __init__(self, vllm_model):
- super().__init__()
self.llm = vllm_model
self.rng = np.random.default_rng(seed=42)
def encode(
self,
- sentences: Sequence[str],
+ inputs: DataLoader[mteb.types.BatchedInput],
*args,
**kwargs,
) -> np.ndarray:
# Hoping to discover potential scheduling
# issues by randomizing the order.
+ sentences = [text for batch in inputs for text in batch["text"]]
r = self.rng.permutation(len(sentences))
sentences = [sentences[i] for i in r]
outputs = self.llm.embed(sentences, use_tqdm=False)
@@ -52,36 +77,70 @@ class VllmMtebEncoder(mteb.Encoder):
embeds = embeds[np.argsort(r)]
return embeds
+ def similarity(
+ self,
+ embeddings1: np.ndarray,
+ embeddings2: np.ndarray,
+ ) -> np.ndarray:
+ # Cosine similarity
+ norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
+ norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
+ sim = np.dot(embeddings1, embeddings2.T) / (norm1 * norm2.T)
+ return sim
+
+ def similarity_pairwise(
+ self,
+ embeddings1: Array,
+ embeddings2: Array,
+ ) -> Array:
+ # Cosine similarity
+ norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
+ norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
+ sim = np.sum(embeddings1 * embeddings2, axis=1) / (
+ norm1.flatten() * norm2.flatten()
+ )
+ return sim
+
+
+class VllmMtebCrossEncoder(mteb.CrossEncoderProtocol):
+ mteb_model_meta = _empty_model_meta
+
+ def __init__(self, vllm_model):
+ self.llm = vllm_model
+ self.rng = np.random.default_rng(seed=42)
+
def predict(
self,
- sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
+ inputs1: DataLoader[mteb.types.BatchedInput],
+ inputs2: DataLoader[mteb.types.BatchedInput],
*args,
**kwargs,
) -> np.ndarray:
- r = self.rng.permutation(len(sentences))
- sentences = [sentences[i] for i in r]
-
- queries = [s[0] for s in sentences]
- corpus = [s[1] for s in sentences]
+ queries = [text for batch in inputs1 for text in batch["text"]]
+ corpus = [text for batch in inputs2 for text in batch["text"]]
outputs = self.llm.score(
queries, corpus, truncate_prompt_tokens=-1, use_tqdm=False
)
scores = np.array(outputs)
- scores = scores[np.argsort(r)]
return scores
-class OpenAIClientMtebEncoder(mteb.Encoder):
+class OpenAIClientMtebEncoder(VllmMtebEncoder):
def __init__(self, model_name: str, client):
- super().__init__()
self.model_name = model_name
self.client = client
self.rng = np.random.default_rng(seed=42)
- def encode(self, sentences: Sequence[str], *args, **kwargs) -> np.ndarray:
+ def encode(
+ self,
+ inputs: DataLoader[mteb.types.BatchedInput],
+ *args,
+ **kwargs,
+ ) -> np.ndarray:
# Hoping to discover potential scheduling
# issues by randomizing the order.
+ sentences = [text for batch in inputs for text in batch["text"]]
r = self.rng.permutation(len(sentences))
sentences = [sentences[i] for i in r]
@@ -94,28 +153,29 @@ class OpenAIClientMtebEncoder(mteb.Encoder):
return embeds
-class ScoreClientMtebEncoder(mteb.Encoder):
+class ScoreClientMtebEncoder(mteb.CrossEncoderProtocol):
+ mteb_model_meta = _empty_model_meta
+
def __init__(self, model_name: str, url):
- super().__init__()
self.model_name = model_name
self.url = url
self.rng = np.random.default_rng(seed=42)
def predict(
self,
- sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
+ inputs1: DataLoader[mteb.types.BatchedInput],
+ inputs2: DataLoader[mteb.types.BatchedInput],
*args,
**kwargs,
) -> np.ndarray:
- r = self.rng.permutation(len(sentences))
- sentences = [sentences[i] for i in r]
+ queries = [text for batch in inputs1 for text in batch["text"]]
+ full_corpus = [text for batch in inputs2 for text in batch["text"]]
outputs = []
- for query, corpus, prompt in sentences:
+ for query, corpus in zip(queries, full_corpus):
outputs.append(self.get_score(query, corpus))
scores = np.array(outputs)
- scores = scores[np.argsort(r)]
return scores
def get_score(self, query, corpus):
@@ -145,16 +205,13 @@ class RerankClientMtebEncoder(ScoreClientMtebEncoder):
return response["results"][0]["relevance_score"]
-def run_mteb_embed_task(encoder, tasks):
+def run_mteb_embed_task(encoder: mteb.EncoderProtocol, tasks):
tasks = mteb.get_tasks(tasks=tasks)
- evaluation = mteb.MTEB(tasks=tasks)
- results = evaluation.run(
+ results = mteb.evaluate(
encoder,
- verbosity=0,
- output_folder=None,
- encode_kwargs={
- "show_progress_bar": False,
- },
+ tasks,
+ cache=None,
+ show_progress_bar=False,
)
main_score = results[0].scores["test"][0]["main_score"]
@@ -244,33 +301,39 @@ def mteb_test_embed_models(
assert st_main_score - vllm_main_score < atol
-def run_mteb_rerank(cross_encoder, tasks, languages):
- with tempfile.TemporaryDirectory() as results_folder:
+def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
+ with tempfile.TemporaryDirectory() as prediction_folder:
bm25s = mteb.get_model("bm25s")
- tasks = mteb.get_tasks(tasks=tasks, languages=languages)
-
- subset = "default"
eval_splits = ["test"]
- evaluation = mteb.MTEB(tasks=tasks)
- evaluation.run(
- bm25s,
- verbosity=0,
- eval_splits=eval_splits,
- save_predictions=True,
- output_folder=f"{results_folder}/stage1",
- encode_kwargs={"show_progress_bar": False},
+ mteb_tasks: list[mteb.abstasks.AbsTaskRetrieval] = mteb.get_tasks(
+ tasks=tasks, languages=languages, eval_splits=eval_splits
)
- results = evaluation.run(
+ mteb.evaluate(
+ bm25s,
+ mteb_tasks,
+ prediction_folder=prediction_folder,
+ show_progress_bar=False,
+ # don't save results for test runs
+ cache=None,
+ overwrite_strategy="always",
+ )
+
+ second_stage_tasks = []
+ for task in mteb_tasks:
+ second_stage_tasks.append(
+ task.convert_to_reranking(
+ prediction_folder,
+ top_k=10,
+ )
+ )
+
+ results = mteb.evaluate(
cross_encoder,
- verbosity=0,
- eval_splits=eval_splits,
- top_k=10,
- save_predictions=True,
- output_folder=f"{results_folder}/stage2",
- previous_results=f"{results_folder}/stage1/NFCorpus_{subset}_predictions.json",
- encode_kwargs={"show_progress_bar": False},
+ second_stage_tasks,
+ show_progress_bar=False,
+ cache=None,
)
main_score = results[0].scores["test"][0]["main_score"]
return main_score
@@ -280,20 +343,6 @@ def mteb_test_rerank_models_hf(
hf_runner, model_name, hf_dtype="float32", hf_model_callback=None
):
with hf_runner(model_name, is_cross_encoder=True, dtype=hf_dtype) as hf_model:
- original_predict = hf_model.predict
-
- def _predict(
- sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
- *args,
- **kwargs,
- ):
- # vllm and st both remove the prompt, fair comparison.
- prompts = [(s[0], s[1]) for s in sentences]
- return original_predict(prompts, *args, **kwargs, batch_size=8)
-
- hf_model.predict = _predict
- hf_model.original_predict = original_predict
-
if hf_model_callback is not None:
hf_model_callback(hf_model)
@@ -310,7 +359,7 @@ def mteb_test_rerank_models(
model_info: RerankModelInfo,
vllm_extra_kwargs=None,
hf_model_callback=None,
- vllm_mteb_encoder=VllmMtebEncoder,
+ vllm_mteb_encoder=VllmMtebCrossEncoder,
atol=MTEB_RERANK_TOL,
):
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
diff --git a/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py b/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
index 2927a37111364..6b2e469644926 100644
--- a/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
+++ b/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
@@ -2,13 +2,15 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
+import mteb
import numpy as np
import pytest
import torch
+from torch.utils.data import DataLoader
from tests.conftest import HfRunner
from tests.models.language.pooling_mteb_test.mteb_utils import (
- VllmMtebEncoder,
+ VllmMtebCrossEncoder,
mteb_test_rerank_models,
)
from tests.models.utils import LASTPoolingRerankModelInfo, RerankModelInfo
@@ -103,7 +105,7 @@ class GemmaRerankerHfRunner(HfRunner):
return torch.Tensor(scores)
-class GemmaMtebEncoder(VllmMtebEncoder):
+class GemmaMtebEncoder(VllmMtebCrossEncoder):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.query_template = "A: {query}\n"
@@ -111,17 +113,26 @@ class GemmaMtebEncoder(VllmMtebEncoder):
def predict(
self,
- sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
+ inputs1: DataLoader[mteb.types.BatchedInput],
+ inputs2: DataLoader[mteb.types.BatchedInput],
*args,
**kwargs,
) -> np.ndarray:
- _sentences = []
- for query, corpus, prompt in sentences:
- query = self.query_template.format(query=query)
- corpus = self.document_template.format(doc=corpus, prompt=PROMPT)
- _sentences.append((query, corpus, prompt))
-
- return super().predict(_sentences, *args, **kwargs)
+ queries = [
+ self.query_template.format(query=text)
+ for batch in inputs1
+ for text in batch["text"]
+ ]
+ corpus = [
+ self.document_template.format(doc=text, prompt=PROMPT)
+ for batch in inputs2
+ for text in batch["text"]
+ ]
+ outputs = self.llm.score(
+ queries, corpus, truncate_prompt_tokens=-1, use_tqdm=False
+ )
+ scores = np.array(outputs)
+ return scores
@pytest.mark.parametrize("model_info", RERANK_MODELS)
diff --git a/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py b/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
index fd04dc1990238..a6f2a89b268f1 100644
--- a/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
+++ b/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
@@ -70,8 +70,9 @@ class MxbaiRerankerHfRunner(HfRunner):
return scores
scores = []
- for prompt in prompts:
- inputs = process_inputs([prompt])
+ for query, doc, *_ in prompts:
+ pairs = [(query, doc)]
+ inputs = process_inputs(pairs)
score = compute_logits(inputs)
scores.append(score[0].item())
return torch.Tensor(scores)
diff --git a/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py b/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
index 00e99f44cfdb1..9a1be6c0be1d6 100644
--- a/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
+++ b/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
@@ -72,8 +72,9 @@ class Qwen3RerankerHfRunner(HfRunner):
return scores
scores = []
- for prompt in prompts:
- inputs = process_inputs([prompt])
+ for query, doc, *_ in prompts:
+ pairs = [(query, doc)]
+ inputs = process_inputs(pairs)
score = compute_logits(inputs)
scores.append(score[0].item())
return torch.Tensor(scores)
diff --git a/tests/models/multimodal/generation/test_multimodal_gguf.py b/tests/models/multimodal/generation/test_multimodal_gguf.py
new file mode 100644
index 0000000000000..e596b20c6302b
--- /dev/null
+++ b/tests/models/multimodal/generation/test_multimodal_gguf.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Literal, NamedTuple
+
+import pytest
+from huggingface_hub import hf_hub_download
+from pytest import MarkDecorator
+
+from tests.quantization.utils import is_quant_method_supported
+from vllm.assets.image import ImageAsset
+from vllm.utils.torch_utils import set_default_torch_num_threads
+
+from ....conftest import PromptImageInput, VllmRunner
+from ...utils import check_logprobs_close
+
+
+class GGUFMMTestConfig(NamedTuple):
+ original_model: str
+ gguf_repo: str
+ gguf_backbone: str
+ gguf_mmproj: str
+ prompt: list[str]
+ mm_data: dict[Literal["images"], PromptImageInput]
+ max_model_len: int = 4096
+ marks: list[MarkDecorator] = []
+
+ @property
+ def gguf_model(self):
+ hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj)
+ return hf_hub_download(self.gguf_repo, filename=self.gguf_backbone)
+
+
+GEMMA3_CONFIG = GGUFMMTestConfig(
+ original_model="google/gemma-3-4b-it",
+ gguf_repo="google/gemma-3-4b-it-qat-q4_0-gguf",
+ gguf_backbone="gemma-3-4b-it-q4_0.gguf",
+ gguf_mmproj="mmproj-model-f16-4B.gguf",
+ prompt=["Describe this image in detail:"],
+ mm_data={"images": [ImageAsset("stop_sign").pil_image]},
+ marks=[pytest.mark.core_model],
+)
+
+MODELS_TO_TEST = [GEMMA3_CONFIG]
+
+
+def run_multimodal_gguf_test(
+ vllm_runner: type[VllmRunner],
+ model: GGUFMMTestConfig,
+ dtype: str,
+ max_tokens: int,
+ num_logprobs: int,
+):
+ # Run gguf model.
+ with (
+ set_default_torch_num_threads(1),
+ vllm_runner(
+ model_name=model.gguf_model,
+ enforce_eager=True,
+ tokenizer_name=model.original_model,
+ dtype=dtype,
+ max_model_len=model.max_model_len,
+ ) as gguf_model,
+ ):
+ gguf_outputs = gguf_model.generate_greedy_logprobs(
+ prompts=model.prompt,
+ max_tokens=max_tokens,
+ num_logprobs=num_logprobs,
+ **model.mm_data,
+ )
+
+ # Run unquantized model.
+ with vllm_runner(
+ model_name=model.original_model,
+ enforce_eager=True, # faster tests
+ dtype=dtype,
+ max_model_len=model.max_model_len,
+ ) as original_model:
+ original_outputs = original_model.generate_greedy_logprobs(
+ prompts=model.prompt,
+ max_tokens=max_tokens,
+ num_logprobs=num_logprobs,
+ **model.mm_data,
+ )
+
+ check_logprobs_close(
+ outputs_0_lst=original_outputs,
+ outputs_1_lst=gguf_outputs,
+ name_0="original",
+ name_1="gguf",
+ )
+
+
+@pytest.mark.skipif(
+ not is_quant_method_supported("gguf"),
+ reason="gguf is not supported on this GPU type.",
+)
+@pytest.mark.parametrize(
+ "model",
+ [
+ pytest.param(test_config, marks=test_config.marks)
+ for test_config in MODELS_TO_TEST
+ ],
+)
+@pytest.mark.parametrize("dtype", ["bfloat16"])
+@pytest.mark.parametrize("max_tokens", [32])
+@pytest.mark.parametrize("num_logprobs", [10])
+def test_models(
+ vllm_runner: type[VllmRunner],
+ model: GGUFMMTestConfig,
+ dtype: str,
+ max_tokens: int,
+ num_logprobs: int,
+) -> None:
+ run_multimodal_gguf_test(vllm_runner, model, dtype, max_tokens, num_logprobs)
diff --git a/tests/models/multimodal/generation/test_qwen2_vl.py b/tests/models/multimodal/generation/test_qwen2_vl.py
index e10b8e1e77af1..e1b7dbf99f1fd 100644
--- a/tests/models/multimodal/generation/test_qwen2_vl.py
+++ b/tests/models/multimodal/generation/test_qwen2_vl.py
@@ -128,12 +128,7 @@ def batch_make_image_embeddings(
visual = model.visual
pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype)
- image_grid_thw_on_device = image_grid_thw.to(
- visual.device, dtype=torch.int64
- )
- return visual(
- pixel_values_on_device, grid_thw=image_grid_thw_on_device
- ).cpu()
+ return visual(pixel_values_on_device, grid_thw=image_grid_thw).cpu()
image_embeds = torch.concat(llm.apply_model(get_image_embeds))
@@ -217,12 +212,7 @@ def batch_make_video_embeddings(
visual = model.visual
pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype)
- video_grid_thw_on_device = video_grid_thw.to(
- visual.device, dtype=torch.int64
- )
- return visual(
- pixel_values_on_device, grid_thw=video_grid_thw_on_device
- ).cpu()
+ return visual(pixel_values_on_device, grid_thw=video_grid_thw).cpu()
video_embeds = torch.concat(llm.apply_model(get_image_embeds))
diff --git a/tests/models/multimodal/test_mapping.py b/tests/models/multimodal/test_mapping.py
index 2f38dc450ef96..0d2eaca95504e 100644
--- a/tests/models/multimodal/test_mapping.py
+++ b/tests/models/multimodal/test_mapping.py
@@ -50,12 +50,24 @@ def test_hf_model_weights_mapper(model_arch: str):
model_info.check_available_online(on_fail="skip")
model_info.check_transformers_version(on_fail="skip")
+ is_mistral_model = model_arch in [
+ "Mistral3ForConditionalGeneration",
+ "PixtralForConditionalGeneration",
+ "VoxtralForConditionalGeneration",
+ ]
+
+ if not is_mistral_model or model_info.tokenizer_mode == "mistral":
+ tokenizer_mode = model_info.tokenizer_mode
+ else:
+ tokenizer_mode = "hf"
+
model_id = model_info.default
model_config = ModelConfig(
model_id,
tokenizer=model_info.tokenizer or model_id,
- tokenizer_mode=model_info.tokenizer_mode,
+ tokenizer_mode=tokenizer_mode,
+ config_format="hf",
revision=model_info.revision,
trust_remote_code=model_info.trust_remote_code,
hf_overrides=model_info.hf_overrides,
diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py
index 24220978534ca..5b8aaa299fdc1 100644
--- a/tests/models/quantization/test_bitsandbytes.py
+++ b/tests/models/quantization/test_bitsandbytes.py
@@ -14,10 +14,13 @@ from vllm.platforms import current_platform
from ...utils import compare_two_settings, multi_gpu_test
from ..utils import check_embeddings_close, check_logprobs_close
-pytestmark = pytest.mark.skipif(
- current_platform.is_rocm(),
- reason="bitsandbytes quantization not supported on ROCm (CUDA-only kernels)",
-)
+if current_platform.is_rocm():
+ from vllm.platforms.rocm import on_gfx9
+
+ pytestmark = pytest.mark.skipif(
+ on_gfx9(),
+ reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)",
+ )
models_4bit_to_test = [
("facebook/opt-125m", "quantize opt model inflight"),
@@ -256,6 +259,9 @@ def validate_generated_texts(
tensor_parallel_size=vllm_tp_size,
enforce_eager=False,
default_torch_num_threads=1,
+ tokenizer_mode="hf",
+ load_format="hf",
+ config_format="hf",
) as llm:
vllm_outputs = llm.generate_greedy(prompts, max_tokens)
vllm_logs = log_generated_texts(prompts, vllm_outputs, "VllmRunner")
diff --git a/tests/models/quantization/test_gguf.py b/tests/models/quantization/test_gguf.py
index 5e2438857aeef..3b9597507ac1b 100644
--- a/tests/models/quantization/test_gguf.py
+++ b/tests/models/quantization/test_gguf.py
@@ -78,6 +78,12 @@ DOLPHIN_CONFIG = GGUFTestConfig(
gguf_filename="tinydolphin-2.8-1.1b.Q6_K.gguf",
)
+GEMMA3_CONFIG = GGUFTestConfig(
+ original_model="google/gemma-3-270m-it",
+ gguf_repo="ggml-org/gemma-3-270m-it-qat-GGUF",
+ gguf_filename="gemma-3-270m-it-qat-Q4_0.gguf",
+)
+
MODELS = [
# LLAMA_CONFIG, # broken: https://github.com/vllm-project/vllm/issues/19458
QWEN2_CONFIG,
@@ -85,6 +91,7 @@ MODELS = [
GPT2_CONFIG,
STABLELM_CONFIG,
DOLPHIN_CONFIG,
+ GEMMA3_CONFIG,
# STARCODER_CONFIG, # broken
]
@@ -148,7 +155,7 @@ def check_model_outputs(
"model",
[pytest.param(test_config, marks=test_config.marks) for test_config in MODELS],
)
-@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.parametrize("dtype", ["bfloat16"])
@pytest.mark.parametrize("max_tokens", [32])
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("tp_size", [1])
diff --git a/tests/models/registry.py b/tests/models/registry.py
index 094f921e4305f..758ec54493aa3 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -370,7 +370,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
),
"OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
"Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
- "Olmo3ForCausalLM": _HfExamplesInfo("shanearora/2025-sep-a-base-model"),
+ "Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"),
"OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),
"OpenPanguMTPModel": _HfExamplesInfo(
"FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",
@@ -402,6 +402,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"pfnet/plamo-2-1b",
trust_remote_code=True,
),
+ "Plamo3ForCausalLM": _HfExamplesInfo(
+ "pfnet/plamo-3-nict-2b-base",
+ trust_remote_code=True,
+ ),
"QWenLMHeadModel": _HfExamplesInfo(
"Qwen/Qwen-7B-Chat",
max_transformers_version="4.53",
@@ -721,6 +725,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"NemotronH_Nano_VL_V2": _HfExamplesInfo(
"nano_vl_dummy", is_available_online=False, trust_remote_code=True
),
+ "OpenCUAForConditionalGeneration": _HfExamplesInfo(
+ "xlangai/OpenCUA-7B", trust_remote_code=True
+ ),
"Ovis": _HfExamplesInfo(
"AIDC-AI/Ovis2-1B",
trust_remote_code=True,
diff --git a/tests/multimodal/assets/corrupted.mp4 b/tests/multimodal/assets/corrupted.mp4
new file mode 100644
index 0000000000000..c355bb932ceee
Binary files /dev/null and b/tests/multimodal/assets/corrupted.mp4 differ
diff --git a/tests/multimodal/test_utils.py b/tests/multimodal/test_utils.py
index ea795fcbbde55..639e290406fe2 100644
--- a/tests/multimodal/test_utils.py
+++ b/tests/multimodal/test_utils.py
@@ -16,10 +16,10 @@ from vllm.multimodal.utils import MediaConnector, argsort_mm_positions
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
- "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
- "Grayscale_8bits_palette_sample_image.png", # "https://upload.wikimedia.org/wikipedia/commons/f/fa/Grayscale_8bits_palette_sample_image.png",
- "1280px-Venn_diagram_rgb.svg.png", # "https://upload.wikimedia.org/wikipedia/commons/thumb/9/91/Venn_diagram_rgb.svg/1280px-Venn_diagram_rgb.svg.png",
- "RGBA_comp.png", # "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
+ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
+ "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",
]
TEST_VIDEO_URLS = [
diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py
index 6572616769a91..6ed21de368ac3 100644
--- a/tests/multimodal/test_video.py
+++ b/tests/multimodal/test_video.py
@@ -18,6 +18,7 @@ from .utils import cosine_similarity, create_video_from_image, normalize_image
pytestmark = pytest.mark.cpu_test
+ASSETS_DIR = Path(__file__).parent / "assets"
NUM_FRAMES = 10
FAKE_OUTPUT_1 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
@@ -140,3 +141,39 @@ def test_opencv_video_io_colorspace(is_color: bool, fourcc: str, ext: str):
)
assert np.sum(np.isnan(sim)) / sim.size < 0.001
assert np.nanmean(sim) > 0.99
+
+
+def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch):
+ """
+ Regression test for handling videos with broken frames.
+ This test uses a pre-corrupted video file (assets/corrupted.mp4) that
+ contains broken/unreadable frames to verify the video loader handles
+ them gracefully without crashing and returns accurate metadata.
+ """
+ with monkeypatch.context() as m:
+ m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
+
+ # Load the pre-corrupted video file that contains broken frames
+ corrupted_video_path = ASSETS_DIR / "corrupted.mp4"
+
+ with open(corrupted_video_path, "rb") as f:
+ video_data = f.read()
+
+ loader = VIDEO_LOADER_REGISTRY.load("opencv")
+ frames, metadata = loader.load_bytes(video_data, num_frames=-1)
+
+ # Verify metadata consistency:
+ # frames_indices must match actual loaded frames
+ assert frames.shape[0] == len(metadata["frames_indices"]), (
+ f"Frames array size must equal frames_indices length. "
+ f"Got {frames.shape[0]} frames but "
+ f"{len(metadata['frames_indices'])} indices"
+ )
+
+ # Verify that broken frames were skipped:
+ # loaded frames should be less than total
+ assert frames.shape[0] < metadata["total_num_frames"], (
+ f"Should load fewer frames than total due to broken frames. "
+ f"Expected fewer than {metadata['total_num_frames']} frames, "
+ f"but loaded {frames.shape[0]} frames"
+ )
diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py
index e7d902ed26aaa..31b65189b5ec3 100644
--- a/tests/quantization/test_compressed_tensors.py
+++ b/tests/quantization/test_compressed_tensors.py
@@ -141,7 +141,7 @@ def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
"neuralmagic/Llama-3.2-1B-quantized.w8a8",
],
)
-@pytest.mark.parametrize("max_tokens", [8])
+@pytest.mark.parametrize("max_tokens", [4])
@pytest.mark.parametrize("num_logprobs", [10])
@pytest.mark.parametrize(
"use_aiter", [True, False] if current_platform.is_rocm() else [False]
@@ -182,7 +182,7 @@ def test_compressed_tensors_w8a8_logprobs(
example_prompts, max_tokens, num_logprobs
)
- with vllm_runner(model_path, dtype=dtype) as vllm_model:
+ with vllm_runner(model_path, dtype=dtype, enforce_eager=True) as vllm_model:
vllm_outputs = vllm_model.generate_greedy_logprobs(
example_prompts, max_tokens, num_logprobs
)
diff --git a/tests/quantization/test_cpu_offload.py b/tests/quantization/test_cpu_offload.py
index a3fb4a6953474..1591ce1c4f5ad 100644
--- a/tests/quantization/test_cpu_offload.py
+++ b/tests/quantization/test_cpu_offload.py
@@ -19,8 +19,8 @@ def test_cpu_offload_fp8():
# Test loading a quantized checkpoint
compare_two_settings(
"neuralmagic/Qwen2-1.5B-Instruct-FP8",
- [],
- ["--cpu-offload-gb", "1"],
+ ["--enforce_eager"],
+ ["--enforce_eager", "--cpu-offload-gb", "1"],
max_wait_seconds=480,
)
@@ -35,8 +35,8 @@ def test_cpu_offload_gptq(monkeypatch):
# Test GPTQ Marlin
compare_two_settings(
"Qwen/Qwen2-1.5B-Instruct-GPTQ-Int4",
- [],
- ["--cpu-offload-gb", "1"],
+ ["--enforce_eager"],
+ ["--enforce_eager", "--cpu-offload-gb", "1"],
max_wait_seconds=480,
)
@@ -51,8 +51,8 @@ def test_cpu_offload_awq(monkeypatch):
# Test AWQ Marlin
compare_two_settings(
"Qwen/Qwen2-1.5B-Instruct-AWQ",
- [],
- ["--cpu-offload-gb", "1"],
+ ["--enforce_eager"],
+ ["--enforce_eager", "--cpu-offload-gb", "1"],
max_wait_seconds=480,
)
@@ -67,7 +67,7 @@ def test_cpu_offload_compressed_tensors(monkeypatch):
# Test wNa16
compare_two_settings(
"nm-testing/tinyllama-oneshot-w4a16-channel-v2",
- [],
- ["--cpu-offload-gb", "1"],
+ ["--enforce_eager"],
+ ["--enforce_eager", "--cpu-offload-gb", "1"],
max_wait_seconds=480,
)
diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py
new file mode 100644
index 0000000000000..077b802e559dc
--- /dev/null
+++ b/tests/quantization/test_cpu_wna16.py
@@ -0,0 +1,23 @@
+# 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 = [
+ "TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
+ "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx
+]
+DTYPE = ["bfloat16"]
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", DTYPE)
+def test_ipex_quant(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/tests/quantization/test_experts_int8.py b/tests/quantization/test_experts_int8.py
index 2a72f734e431b..b992e976ac308 100644
--- a/tests/quantization/test_experts_int8.py
+++ b/tests/quantization/test_experts_int8.py
@@ -21,7 +21,7 @@ MODELS = ["ai21labs/Jamba-tiny-random", "pfnet/plamo-2-1b"]
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["bfloat16"])
-@pytest.mark.parametrize("max_tokens", [10])
+@pytest.mark.parametrize("max_tokens", [4])
def test_model_experts_int8_startup(
hf_runner,
vllm_runner,
@@ -33,5 +33,7 @@ def test_model_experts_int8_startup(
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
model_info.check_transformers_version(on_fail="skip")
- with vllm_runner(model, dtype=dtype, quantization="experts_int8") as vllm_model:
+ with vllm_runner(
+ model, dtype=dtype, enforce_eager=True, quantization="experts_int8"
+ ) as vllm_model:
vllm_model.generate_greedy(example_prompts, max_tokens)
diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py
index f02da2996ffea..7bcac9ad768e7 100644
--- a/tests/quantization/test_fp8.py
+++ b/tests/quantization/test_fp8.py
@@ -45,10 +45,10 @@ def test_model_load_and_run(
if force_marlin:
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
- with vllm_runner(model_id) as llm:
+ with vllm_runner(model_id, enforce_eager=True) as llm:
# note: this does not test accuracy, just that we can run through
# see lm-eval tests for accuracy
- outputs = llm.generate_greedy(["Hello my name is"], max_tokens=10)
+ outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
print(outputs[0][1])
@@ -85,7 +85,7 @@ def test_kv_cache_model_load_and_run(
# `LLM.apply_model` requires pickling a function.
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
- with vllm_runner(model_id, kv_cache_dtype="fp8") as llm:
+ with vllm_runner(model_id, kv_cache_dtype="fp8", enforce_eager=True) as llm:
def check_model(model):
attn = model.model.layers[0].self_attn.attn
@@ -112,7 +112,7 @@ def test_kv_cache_model_load_and_run(
# note: this does not test accuracy, just that we can run through
# see lm-eval tests for accuracy
- outputs = llm.generate_greedy(["Hello my name is"], max_tokens=10)
+ outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
print(outputs[0][1])
@@ -142,7 +142,10 @@ def test_load_fp16_model(
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
with vllm_runner(
- "facebook/opt-125m", quantization="fp8", kv_cache_dtype=kv_cache_dtype
+ "facebook/opt-125m",
+ quantization="fp8",
+ enforce_eager=True,
+ kv_cache_dtype=kv_cache_dtype,
) as llm:
def check_model(model):
diff --git a/tests/quantization/test_ipex_quant.py b/tests/quantization/test_ipex_quant.py
index ae9b1df3377dc..4f3c52df6c283 100644
--- a/tests/quantization/test_ipex_quant.py
+++ b/tests/quantization/test_ipex_quant.py
@@ -26,7 +26,7 @@ DTYPE = ["bfloat16"]
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", DTYPE)
def test_ipex_quant(vllm_runner, model, dtype):
- with vllm_runner(model, dtype=dtype) as llm:
- output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
+ with vllm_runner(model, dtype=dtype, enforce_eager=True) as llm:
+ output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
assert output
print(output)
diff --git a/tests/quantization/test_lm_head.py b/tests/quantization/test_lm_head.py
index f009a4cfb870d..d92dfaa2cc7b5 100644
--- a/tests/quantization/test_lm_head.py
+++ b/tests/quantization/test_lm_head.py
@@ -49,4 +49,4 @@ def test_lm_head(
vllm_model.apply_model(check_model)
- print(vllm_model.generate_greedy(["Hello my name is"], max_tokens=10)[0][1])
+ print(vllm_model.generate_greedy(["Hello my name is"], max_tokens=4)[0][1])
diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py
index 8abf65d29784d..0298994c396f6 100644
--- a/tests/quantization/test_modelopt.py
+++ b/tests/quantization/test_modelopt.py
@@ -88,6 +88,6 @@ def test_modelopt_fp8_checkpoint_setup(vllm_runner):
llm.apply_model(check_model)
# Run a simple generation test to ensure the model works
- output = llm.generate_greedy(["Hello my name is"], max_tokens=20)
+ output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
assert output
print(f"ModelOpt FP8 output: {output}")
diff --git a/tests/quantization/test_ptpc_fp8.py b/tests/quantization/test_ptpc_fp8.py
index e8ea4148585bf..61efd2ce66c71 100644
--- a/tests/quantization/test_ptpc_fp8.py
+++ b/tests/quantization/test_ptpc_fp8.py
@@ -38,6 +38,7 @@ def test_ptpc_fp8_rocm(vllm_runner, dtype: str, kv_cache_dtype: str) -> None:
"facebook/opt-125m",
dtype=dtype,
quantization="ptpc_fp8",
+ enforce_eager=True,
kv_cache_dtype=kv_cache_dtype,
)
except AssertionError as e:
@@ -65,5 +66,5 @@ def test_ptpc_fp8_rocm(vllm_runner, dtype: str, kv_cache_dtype: str) -> None:
llm.apply_model(check_model)
- output = llm.generate_greedy("Hello my name is", max_tokens=20)
+ output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
diff --git a/tests/quantization/test_register_quantization_config.py b/tests/quantization/test_register_quantization_config.py
index 8da048703df93..a09856c78559a 100644
--- a/tests/quantization/test_register_quantization_config.py
+++ b/tests/quantization/test_register_quantization_config.py
@@ -23,8 +23,8 @@ from vllm.model_executor.layers.quantization import (
get_quantization_config,
register_quantization_config,
)
-from vllm.model_executor.layers.quantization.base_config import ( # noqa: E501
- QuantizationConfig,
+from vllm.model_executor.layers.quantization.base_config import (
+ QuantizationConfig, # noqa: E501
)
@@ -142,5 +142,5 @@ def test_custom_quant(vllm_runner, model, monkeypatch):
llm.apply_model(check_model)
- output = llm.generate_greedy("Hello my name is", max_tokens=20)
+ output = llm.generate_greedy("Hello my name is", max_tokens=1)
assert output
diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py
index 82413f36e997f..f35c3973ab6e6 100644
--- a/tests/quantization/test_torchao.py
+++ b/tests/quantization/test_torchao.py
@@ -225,13 +225,12 @@ def test_reload_weights():
@pytest.mark.skip(
reason="since torchao nightly is only compatible with torch nightly"
"currently https://github.com/pytorch/ao/issues/2919, we'll have to skip "
- "torchao tests that requires newer versions (0.14.0.dev+) for now"
+ "torchao tests that requires newer versions (0.15.0.dev+) for now"
)
-def test_opt_125m_float8_weight_only_safetensors_model_loading_with_params(vllm_runner):
+def test_safetensors_model_loading_with_params(vllm_runner):
torch._dynamo.reset()
- model_name = (
- "torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.14.0.dev-safetensors"
- )
+ # using this model to test safetensors loading with file sharding
+ model_name = "torchao-testing/Qwen3-8B-INT4-0.15.0dev-safetensors"
with vllm_runner(model_name=model_name, dtype="bfloat16") as llm:
output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
@@ -392,7 +391,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant(
assert not has_int4_preshuffled_tensor
assert weight_attrs == [False, 1, 0, True]
- output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
+ output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
assert output
diff --git a/tests/samplers/test_logprobs.py b/tests/samplers/test_logprobs.py
index c9d227599cde5..ea40c48027205 100644
--- a/tests/samplers/test_logprobs.py
+++ b/tests/samplers/test_logprobs.py
@@ -24,9 +24,7 @@ def test_ranks(
greedy,
flat_logprobs,
example_prompts,
- monkeypatch: pytest.MonkeyPatch,
):
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1" if flat_logprobs else "0")
with vllm_runner(model, dtype=dtype, max_logprobs=MAX_LOGPROBS) as vllm_model:
tokenizer = vllm_model.llm.get_tokenizer()
example_prompt_tokens = [tokenizer.encode(prompt) for prompt in example_prompts]
@@ -36,6 +34,7 @@ def test_ranks(
max_tokens=MAX_TOKENS,
logprobs=NUM_TOP_LOGPROBS,
prompt_logprobs=NUM_PROMPT_LOGPROBS,
+ flat_logprobs=flat_logprobs,
)
results = vllm_model.generate_w_logprobs(example_prompts, sampling_params)
diff --git a/tests/test_config.py b/tests/test_config.py
index bba2fbec3db29..16f68d18fc68b 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -249,45 +249,48 @@ def test_get_bert_tokenization_sentence_transformer_config():
def test_rope_customization():
- TEST_ROPE_SCALING = {"rope_type": "dynamic", "factor": 2.0}
- TEST_ROPE_THETA = 16_000_000.0
- LONGCHAT_ROPE_SCALING = {"rope_type": "linear", "factor": 8.0}
+ TEST_ROPE_PARAMETERS = {
+ "rope_theta": 16_000_000.0,
+ "rope_type": "dynamic",
+ "factor": 2.0,
+ }
+ LLAMA_ROPE_PARAMETERS = {"rope_theta": 500000.0, "rope_type": "default"}
+ LONGCHAT_ROPE_PARAMETERS = {"rope_type": "linear", "factor": 8.0}
llama_model_config = ModelConfig("meta-llama/Meta-Llama-3-8B-Instruct")
- assert getattr(llama_model_config.hf_config, "rope_scaling", None) is None
- assert getattr(llama_model_config.hf_config, "rope_theta", None) == 500_000
+ assert (
+ getattr(llama_model_config.hf_config, "rope_parameters", None)
+ == LLAMA_ROPE_PARAMETERS
+ )
assert llama_model_config.max_model_len == 8192
llama_model_config = ModelConfig(
"meta-llama/Meta-Llama-3-8B-Instruct",
- hf_overrides={
- "rope_scaling": TEST_ROPE_SCALING,
- "rope_theta": TEST_ROPE_THETA,
- },
+ hf_overrides={"rope_parameters": TEST_ROPE_PARAMETERS},
)
assert (
- getattr(llama_model_config.hf_config, "rope_scaling", None) == TEST_ROPE_SCALING
+ getattr(llama_model_config.hf_config, "rope_parameters", None)
+ == TEST_ROPE_PARAMETERS
)
- assert getattr(llama_model_config.hf_config, "rope_theta", None) == TEST_ROPE_THETA
assert llama_model_config.max_model_len == 16384
longchat_model_config = ModelConfig("lmsys/longchat-13b-16k")
- # Check if LONGCHAT_ROPE_SCALING entries are in longchat_model_config
+ # Check if LONGCHAT_ROPE_PARAMETERS entries are in longchat_model_config
assert all(
- longchat_model_config.hf_config.rope_scaling.get(key) == value
- for key, value in LONGCHAT_ROPE_SCALING.items()
+ longchat_model_config.hf_config.rope_parameters.get(key) == value
+ for key, value in LONGCHAT_ROPE_PARAMETERS.items()
)
assert longchat_model_config.max_model_len == 16384
longchat_model_config = ModelConfig(
"lmsys/longchat-13b-16k",
hf_overrides={
- "rope_scaling": TEST_ROPE_SCALING,
+ "rope_parameters": TEST_ROPE_PARAMETERS,
},
)
assert (
- getattr(longchat_model_config.hf_config, "rope_scaling", None)
- == TEST_ROPE_SCALING
+ getattr(longchat_model_config.hf_config, "rope_parameters", None)
+ == TEST_ROPE_PARAMETERS
)
assert longchat_model_config.max_model_len == 4096
diff --git a/tests/test_logger.py b/tests/test_logger.py
index 01672358902f9..8900e9c2a1e69 100644
--- a/tests/test_logger.py
+++ b/tests/test_logger.py
@@ -49,10 +49,13 @@ def test_trace_function_call():
os.remove(path)
-def test_default_vllm_root_logger_configuration():
+def test_default_vllm_root_logger_configuration(monkeypatch):
"""This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and
VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default
behavior is activated."""
+ monkeypatch.setenv("VLLM_LOGGING_COLOR", "0")
+ _configure_vllm_root_logger()
+
logger = logging.getLogger("vllm")
assert logger.level == logging.DEBUG
assert not logger.propagate
@@ -70,12 +73,13 @@ def test_default_vllm_root_logger_configuration():
assert formatter.datefmt == _DATE_FORMAT
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 1)
-@patch("vllm.logger.VLLM_LOGGING_CONFIG_PATH", None)
-def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger():
+def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypatch):
"""This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and
VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default
behavior is activated."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1")
+ monkeypatch.delenv("VLLM_LOGGING_CONFIG_PATH", raising=False)
+
root_logger = logging.getLogger("vllm")
root_handler = root_logger.handlers[0]
@@ -99,49 +103,50 @@ def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger():
assert log_record.levelno == logging.INFO
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 0)
-@patch("vllm.logger.VLLM_LOGGING_CONFIG_PATH", None)
-def test_logger_configuring_can_be_disabled():
+def test_logger_configuring_can_be_disabled(monkeypatch):
"""This test calls _configure_vllm_root_logger again to test custom logging
config behavior, however mocks are used to ensure no changes in behavior or
configuration occur."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "0")
+ monkeypatch.delenv("VLLM_LOGGING_CONFIG_PATH", raising=False)
with patch("vllm.logger.dictConfig") as dict_config_mock:
_configure_vllm_root_logger()
dict_config_mock.assert_not_called()
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 1)
-@patch(
- "vllm.logger.VLLM_LOGGING_CONFIG_PATH",
- "/if/there/is/a/file/here/then/you/did/this/to/yourself.json",
-)
-def test_an_error_is_raised_when_custom_logging_config_file_does_not_exist():
+def test_an_error_is_raised_when_custom_logging_config_file_does_not_exist(monkeypatch):
"""This test calls _configure_vllm_root_logger again to test custom logging
config behavior, however it fails before any change in behavior or
configuration occurs."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1")
+ monkeypatch.setenv(
+ "VLLM_LOGGING_CONFIG_PATH",
+ "/if/there/is/a/file/here/then/you/did/this/to/yourself.json",
+ )
+
with pytest.raises(RuntimeError) as ex_info:
_configure_vllm_root_logger()
assert ex_info.type == RuntimeError # noqa: E721
assert "File does not exist" in str(ex_info)
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 1)
-def test_an_error_is_raised_when_custom_logging_config_is_invalid_json():
+def test_an_error_is_raised_when_custom_logging_config_is_invalid_json(monkeypatch):
"""This test calls _configure_vllm_root_logger again to test custom logging
config behavior, however it fails before any change in behavior or
configuration occurs."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1")
+
with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file:
logging_config_file.write("---\nloggers: []\nversion: 1")
logging_config_file.flush()
- with patch("vllm.logger.VLLM_LOGGING_CONFIG_PATH", logging_config_file.name):
- with pytest.raises(JSONDecodeError) as ex_info:
- _configure_vllm_root_logger()
- assert ex_info.type == JSONDecodeError
- assert "Expecting value" in str(ex_info)
+ monkeypatch.setenv("VLLM_LOGGING_CONFIG_PATH", logging_config_file.name)
+ with pytest.raises(JSONDecodeError) as ex_info:
+ _configure_vllm_root_logger()
+ assert ex_info.type == JSONDecodeError
+ assert "Expecting value" in str(ex_info)
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 1)
@pytest.mark.parametrize(
"unexpected_config",
(
@@ -151,26 +156,30 @@ def test_an_error_is_raised_when_custom_logging_config_is_invalid_json():
),
)
def test_an_error_is_raised_when_custom_logging_config_is_unexpected_json(
+ monkeypatch,
unexpected_config: Any,
):
"""This test calls _configure_vllm_root_logger again to test custom logging
config behavior, however it fails before any change in behavior or
configuration occurs."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1")
+
with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file:
logging_config_file.write(json.dumps(unexpected_config))
logging_config_file.flush()
- with patch("vllm.logger.VLLM_LOGGING_CONFIG_PATH", logging_config_file.name):
- with pytest.raises(ValueError) as ex_info:
- _configure_vllm_root_logger()
- assert ex_info.type == ValueError # noqa: E721
- assert "Invalid logging config. Expected dict, got" in str(ex_info)
+ monkeypatch.setenv("VLLM_LOGGING_CONFIG_PATH", logging_config_file.name)
+ with pytest.raises(ValueError) as ex_info:
+ _configure_vllm_root_logger()
+ assert ex_info.type == ValueError # noqa: E721
+ assert "Invalid logging config. Expected dict, got" in str(ex_info)
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 1)
-def test_custom_logging_config_is_parsed_and_used_when_provided():
+def test_custom_logging_config_is_parsed_and_used_when_provided(monkeypatch):
"""This test calls _configure_vllm_root_logger again to test custom logging
config behavior, however mocks are used to ensure no changes in behavior or
configuration occur."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "1")
+
valid_logging_config = {
"loggers": {
"vllm.test_logger.logger": {
@@ -183,19 +192,18 @@ def test_custom_logging_config_is_parsed_and_used_when_provided():
with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file:
logging_config_file.write(json.dumps(valid_logging_config))
logging_config_file.flush()
- with (
- patch("vllm.logger.VLLM_LOGGING_CONFIG_PATH", logging_config_file.name),
- patch("vllm.logger.dictConfig") as dict_config_mock,
- ):
+ monkeypatch.setenv("VLLM_LOGGING_CONFIG_PATH", logging_config_file.name)
+ with patch("vllm.logger.dictConfig") as dict_config_mock:
_configure_vllm_root_logger()
dict_config_mock.assert_called_with(valid_logging_config)
-@patch("vllm.logger.VLLM_CONFIGURE_LOGGING", 0)
-def test_custom_logging_config_causes_an_error_if_configure_logging_is_off():
+def test_custom_logging_config_causes_an_error_if_configure_logging_is_off(monkeypatch):
"""This test calls _configure_vllm_root_logger again to test custom logging
config behavior, however mocks are used to ensure no changes in behavior or
configuration occur."""
+ monkeypatch.setenv("VLLM_CONFIGURE_LOGGING", "0")
+
valid_logging_config = {
"loggers": {
"vllm.test_logger.logger": {
@@ -207,15 +215,15 @@ def test_custom_logging_config_causes_an_error_if_configure_logging_is_off():
with NamedTemporaryFile(encoding="utf-8", mode="w") as logging_config_file:
logging_config_file.write(json.dumps(valid_logging_config))
logging_config_file.flush()
- with patch("vllm.logger.VLLM_LOGGING_CONFIG_PATH", logging_config_file.name):
- with pytest.raises(RuntimeError) as ex_info:
- _configure_vllm_root_logger()
- assert ex_info.type is RuntimeError
- expected_message_snippet = (
- "VLLM_CONFIGURE_LOGGING evaluated to false, but "
- "VLLM_LOGGING_CONFIG_PATH was given."
- )
- assert expected_message_snippet in str(ex_info)
+ monkeypatch.setenv("VLLM_LOGGING_CONFIG_PATH", logging_config_file.name)
+ with pytest.raises(RuntimeError) as ex_info:
+ _configure_vllm_root_logger()
+ assert ex_info.type is RuntimeError
+ expected_message_snippet = (
+ "VLLM_CONFIGURE_LOGGING evaluated to false, but "
+ "VLLM_LOGGING_CONFIG_PATH was given."
+ )
+ assert expected_message_snippet in str(ex_info)
# Remember! The root logger is assumed to have been configured as
# though VLLM_CONFIGURE_LOGGING=1 and VLLM_LOGGING_CONFIG_PATH=None.
diff --git a/tests/test_logprobs.py b/tests/test_logprobs.py
index d26a460d2bcab..75e9d337aa24e 100644
--- a/tests/test_logprobs.py
+++ b/tests/test_logprobs.py
@@ -2,8 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import pytest
-
from vllm.logprobs import (
FlatLogprobs,
Logprob,
@@ -14,24 +12,20 @@ from vllm.logprobs import (
)
-def test_create_logprobs_non_flat(monkeypatch: pytest.MonkeyPatch) -> None:
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "0")
-
- prompt_logprobs = create_prompt_logprobs()
+def test_create_logprobs_non_flat() -> None:
+ prompt_logprobs = create_prompt_logprobs(flat_logprobs=False)
assert isinstance(prompt_logprobs, list)
# Ensure first prompt position logprobs is None
assert len(prompt_logprobs) == 1
assert prompt_logprobs[0] is None
- sample_logprobs = create_sample_logprobs()
+ sample_logprobs = create_sample_logprobs(flat_logprobs=False)
assert isinstance(sample_logprobs, list)
assert len(sample_logprobs) == 0
-def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1")
-
- prompt_logprobs = create_prompt_logprobs()
+def test_create_logprobs_flat() -> None:
+ prompt_logprobs = create_prompt_logprobs(flat_logprobs=True)
assert isinstance(prompt_logprobs, FlatLogprobs)
assert prompt_logprobs.start_indices == [0]
assert prompt_logprobs.end_indices == [0]
@@ -43,7 +37,7 @@ def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
assert len(prompt_logprobs) == 1
assert prompt_logprobs[0] == dict()
- sample_logprobs = create_sample_logprobs()
+ sample_logprobs = create_sample_logprobs(flat_logprobs=True)
assert isinstance(sample_logprobs, FlatLogprobs)
assert len(sample_logprobs.start_indices) == 0
assert len(sample_logprobs.end_indices) == 0
@@ -54,11 +48,8 @@ def test_create_logprobs_flat(monkeypatch: pytest.MonkeyPatch) -> None:
assert len(sample_logprobs) == 0
-def test_append_logprobs_for_next_position_none_flat(
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "0")
- logprobs = create_sample_logprobs()
+def test_append_logprobs_for_next_position_none_flat() -> None:
+ logprobs = create_sample_logprobs(flat_logprobs=False)
append_logprobs_for_next_position(
logprobs,
token_ids=[1],
@@ -85,11 +76,8 @@ def test_append_logprobs_for_next_position_none_flat(
]
-def test_append_logprobs_for_next_position_flat(
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- monkeypatch.setenv("VLLM_FLAT_LOGPROBS", "1")
- logprobs = create_sample_logprobs()
+def test_append_logprobs_for_next_position_flat() -> None:
+ logprobs = create_sample_logprobs(flat_logprobs=True)
append_logprobs_for_next_position(
logprobs,
token_ids=[1],
diff --git a/tests/tool_use/test_kimi_k2_tool_parser.py b/tests/tool_use/test_kimi_k2_tool_parser.py
index 33dabbc7e7b91..3a48b5206141d 100644
--- a/tests/tool_use/test_kimi_k2_tool_parser.py
+++ b/tests/tool_use/test_kimi_k2_tool_parser.py
@@ -60,6 +60,11 @@ def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):
ids=[
"tool_call_with_content_before",
"multi_tool_call_with_content_before",
+ "concatenated_tool_calls_bug_fix",
+ "three_concatenated_tool_calls",
+ "mixed_spacing_tool_calls",
+ "angle_brackets_in_json",
+ "newlines_in_json",
],
argnames=["model_output", "expected_tool_calls", "expected_content"],
argvalues=[
@@ -114,6 +119,123 @@ functions.get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool
],
"I'll help you check the weather. ",
),
+ (
+ """I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"latitude": 34.0522, "longitude": -118.2437}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"content": "Los Angeles today"}<|tool_call_end|><|tool_calls_section_end|>""",
+ [
+ ToolCall(
+ id="functions.get_weather:0",
+ function=FunctionCall(
+ name="get_weather",
+ arguments=json.dumps(
+ {"latitude": 34.0522, "longitude": -118.2437}
+ ),
+ ),
+ type="function",
+ ),
+ ToolCall(
+ id="functions.get_news:1",
+ function=FunctionCall(
+ name="get_news",
+ arguments=json.dumps({"content": "Los Angeles today"}),
+ ),
+ type="function",
+ ),
+ ],
+ "I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. ",
+ ),
+ (
+ """I'll help you with multiple tasks. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "New York"}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"topic": "technology"}<|tool_call_end|><|tool_call_begin|>functions.send_email:2<|tool_call_argument_begin|>{"to": "user@example.com", "subject": "Daily Update"}<|tool_call_end|><|tool_calls_section_end|>""",
+ [
+ ToolCall(
+ id="functions.get_weather:0",
+ function=FunctionCall(
+ name="get_weather",
+ arguments=json.dumps({"city": "New York"}),
+ ),
+ type="function",
+ ),
+ ToolCall(
+ id="functions.get_news:1",
+ function=FunctionCall(
+ name="get_news",
+ arguments=json.dumps({"topic": "technology"}),
+ ),
+ type="function",
+ ),
+ ToolCall(
+ id="functions.send_email:2",
+ function=FunctionCall(
+ name="send_email",
+ arguments=json.dumps(
+ {"to": "user@example.com", "subject": "Daily Update"}
+ ),
+ ),
+ type="function",
+ ),
+ ],
+ "I'll help you with multiple tasks. ",
+ ),
+ (
+ """Mixed spacing test. <|tool_calls_section_begin|> <|tool_call_begin|> functions.test:0 <|tool_call_argument_begin|> {} <|tool_call_end|><|tool_call_begin|>functions.test2:1<|tool_call_argument_begin|>{}<|tool_call_end|> <|tool_calls_section_end|>""",
+ [
+ ToolCall(
+ id="functions.test:0",
+ function=FunctionCall(
+ name="test",
+ arguments=json.dumps({}),
+ ),
+ type="function",
+ ),
+ ToolCall(
+ id="functions.test2:1",
+ function=FunctionCall(
+ name="test2",
+ arguments=json.dumps({}),
+ ),
+ type="function",
+ ),
+ ],
+ "Mixed spacing test. ",
+ ),
+ (
+ """I need to process HTML content. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_html:0<|tool_call_argument_begin|>{"html": "content
", "text": "normal text"}<|tool_call_end|><|tool_calls_section_end|>""",
+ [
+ ToolCall(
+ id="functions.process_html:0",
+ function=FunctionCall(
+ name="process_html",
+ arguments=json.dumps(
+ {"html": "content
", "text": "normal text"}
+ ),
+ ),
+ type="function",
+ )
+ ],
+ "I need to process HTML content. ",
+ ),
+ (
+ """I need to process formatted JSON. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_data:0<|tool_call_argument_begin|>{
+ "name": "test",
+ "value": 123,
+ "nested": {
+ "key": "value"
+ }
+}<|tool_call_end|><|tool_calls_section_end|>""",
+ [
+ ToolCall(
+ id="functions.process_data:0",
+ function=FunctionCall(
+ name="process_data",
+ arguments=json.dumps(
+ {"name": "test", "value": 123, "nested": {"key": "value"}},
+ indent=2,
+ ),
+ ),
+ type="function",
+ )
+ ],
+ "I need to process formatted JSON. ",
+ ),
],
)
def test_extract_tool_calls(
diff --git a/tests/tool_use/utils.py b/tests/tool_use/utils.py
index 38def6f874d7d..7584b903156b7 100644
--- a/tests/tool_use/utils.py
+++ b/tests/tool_use/utils.py
@@ -128,6 +128,12 @@ CONFIGS: dict[str, ServerConfig] = {
"arguments": [
"--enforce-eager",
"--no-enable-prefix-caching",
+ "--tokenizer_mode",
+ "hf",
+ "--load_format",
+ "hf",
+ "--config_format",
+ "hf",
"--tool-call-parser",
"mistral",
"--chat-template",
@@ -140,21 +146,22 @@ CONFIGS: dict[str, ServerConfig] = {
"without calling a tool. DO NOT CALL A TOOL THAT IS IRRELEVANT "
"to the user's question - just respond to it normally.",
},
- # V1 Test: Passing locally but failing in CI. This runs the
- # V0 Engine because of CPU offloading. Need to debug why.
+ # FIXME: This test currently fails, need to debug why.
# "granite20b": {
- # "model":
- # "mbayser/granite-20b-functioncalling-FP8-KV",
+ # "model": "mbayser/granite-20b-functioncalling-FP8-KV",
# "arguments": [
- # "--tool-call-parser", "granite-20b-fc", "--chat-template",
- # str(VLLM_PATH /
- # "examples/tool_chat_template_granite_20b_fc.jinja"),
- # "--max_num_seqs", "1", "--enforce-eager", "--cpu-offload-gb", "20"
+ # "--tool-call-parser",
+ # "granite-20b-fc",
+ # "--chat-template",
+ # str(VLLM_PATH / "examples/tool_chat_template_granite_20b_fc.jinja"),
+ # "--max_num_seqs",
+ # "1",
+ # "--enforce-eager",
+ # "--cpu-offload-gb",
+ # "20",
# ],
- # "supports_parallel":
- # False,
- # "supports_rocm":
- # False,
+ # "supports_parallel": False,
+ # "supports_rocm": False,
# },
"granite-3.0-8b": {
"model": "ibm-granite/granite-3.0-8b-instruct",
diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py
new file mode 100644
index 0000000000000..de28ab5f99e8c
--- /dev/null
+++ b/tests/transformers_utils/test_config.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, call, patch
+
+import pytest
+
+from vllm.transformers_utils.config import list_filtered_repo_files
+
+
+@pytest.mark.parametrize(
+ "allow_patterns,expected_relative_files",
+ [
+ (
+ ["*.json", "correct*.txt"],
+ ["json_file.json", "subfolder/correct.txt", "correct_2.txt"],
+ ),
+ ],
+)
+def test_list_filtered_repo_files(
+ allow_patterns: list[str], expected_relative_files: list[str]
+):
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ # Prep folder and files
+ path_tmp_dir = Path(tmp_dir)
+ subfolder = path_tmp_dir / "subfolder"
+ subfolder.mkdir()
+ (path_tmp_dir / "json_file.json").touch()
+ (path_tmp_dir / "correct_2.txt").touch()
+ (path_tmp_dir / "uncorrect.txt").touch()
+ (path_tmp_dir / "uncorrect.jpeg").touch()
+ (subfolder / "correct.txt").touch()
+ (subfolder / "uncorrect_sub.txt").touch()
+
+ def _glob_path() -> list[str]:
+ return [
+ str(file.relative_to(path_tmp_dir))
+ for file in path_tmp_dir.glob("**/*")
+ if file.is_file()
+ ]
+
+ # Patch list_repo_files called by fn
+ with patch(
+ "vllm.transformers_utils.config.list_repo_files",
+ MagicMock(return_value=_glob_path()),
+ ) as mock_list_repo_files:
+ out_files = sorted(
+ list_filtered_repo_files(
+ tmp_dir, allow_patterns, "revision", "model", "token"
+ )
+ )
+ assert out_files == sorted(expected_relative_files)
+ assert mock_list_repo_files.call_count == 1
+ assert mock_list_repo_files.call_args_list[0] == call(
+ repo_id=tmp_dir,
+ revision="revision",
+ repo_type="model",
+ token="token",
+ )
diff --git a/tests/transformers_utils/test_utils.py b/tests/transformers_utils/test_utils.py
index beaef04d766bf..bfe1cec76c138 100644
--- a/tests/transformers_utils/test_utils.py
+++ b/tests/transformers_utils/test_utils.py
@@ -2,7 +2,11 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from vllm.transformers_utils.utils import is_cloud_storage, is_gcs, is_s3
+from vllm.transformers_utils.utils import (
+ is_cloud_storage,
+ is_gcs,
+ is_s3,
+)
def test_is_gcs():
diff --git a/tests/utils.py b/tests/utils.py
index 848c4efa8bcde..a87442afb67d5 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -683,7 +683,7 @@ def compare_all_settings(
results += _test_image_text(
client,
model,
- "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
+ "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
)
elif method == "encode":
results += _test_embeddings(client, model, prompt)
diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py
index 1bd05e6183dc2..783e02ce89bdb 100644
--- a/tests/v1/attention/test_mla_backends.py
+++ b/tests/v1/attention/test_mla_backends.py
@@ -61,7 +61,7 @@ for backend in BACKENDS_TO_TEST:
BACKEND_BLOCK_SIZES = {}
for backend in BACKENDS_TO_TEST:
- supported_sizes = backend.get_class().supported_kernel_block_sizes
+ supported_sizes = backend.get_class().get_supported_kernel_block_sizes()
if supported_sizes:
default_size = supported_sizes[0]
block_size = (
diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py
index 1d80ee9875913..e0645ed43015e 100644
--- a/tests/v1/core/test_async_scheduler.py
+++ b/tests/v1/core/test_async_scheduler.py
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections import deque
-import numpy as np
import pytest
from vllm.v1.core.sched.output import SchedulerOutput
@@ -22,7 +21,7 @@ def _make_model_runner_output(
return ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
- sampled_token_ids=[np.array([i]) for i in range(len(req_ids))],
+ sampled_token_ids=[[i] for i in range(len(req_ids))],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
diff --git a/tests/v1/core/test_priority_scheduler_random.py b/tests/v1/core/test_priority_scheduler_random.py
index ba0b703302e38..b4805be802723 100644
--- a/tests/v1/core/test_priority_scheduler_random.py
+++ b/tests/v1/core/test_priority_scheduler_random.py
@@ -3,7 +3,6 @@
import random
import uuid
-import numpy as np
import pytest
from vllm.config import VllmConfig
@@ -100,7 +99,8 @@ def _mock_execute_model(
random.randint(*num_output_tokens_range) for _ in range(len(request_ids))
]
sampled_token_ids = [
- np.random.randint(0, 100, size=num_tokens) for num_tokens in num_output_tokens
+ [random.randint(0, 100) for _ in range(num_tokens)]
+ for num_tokens in num_output_tokens
]
return ModelRunnerOutput(
@@ -196,8 +196,6 @@ def test_priority_scheduling_blast(
num_blocks: int,
):
random.seed(42)
- np.random.seed(42)
-
seen_request_prompt_length = dict[str, int]()
seen_request_ids = set[str]()
seen_mm_hashes = set[str]()
diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 0570c0854c678..09acde6e08faa 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -3,7 +3,6 @@
import dataclasses
from unittest.mock import Mock
-import numpy as np
import pytest
import torch
@@ -77,11 +76,11 @@ def test_get_num_unfinished_requests():
@pytest.mark.parametrize(
"enable_prefix_caching, prompt_logprobs",
[
- (None, None),
+ (False, None),
(True, 5),
],
)
-def test_schedule(enable_prefix_caching: bool | None, prompt_logprobs: int | None):
+def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None):
"""Test scheduling.
Two cases: default APC/no prompt logprobs; APC=True + prompt logprobs
"""
@@ -170,7 +169,7 @@ def test_schedule_partial_requests():
req_id_to_index=req_to_index,
# Only the first request has a sampled token id because
# the rest requests are still being prefilled.
- sampled_token_ids=[np.array([0]), np.array([]), np.array([])],
+ sampled_token_ids=[[0], [], []],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -217,7 +216,7 @@ def test_no_mm_input_chunking():
model_runner_output = ModelRunnerOutput(
req_ids=[request.request_id for request in requests],
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([]) for _ in range(len(requests))],
+ sampled_token_ids=[[] for _ in range(len(requests))],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -277,7 +276,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
model_runner_output = ModelRunnerOutput(
req_ids=[request.request_id for request in requests],
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([]) for _ in range(len(requests))],
+ sampled_token_ids=[[] for _ in range(len(requests))],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -301,8 +300,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
model_runner_output = ModelRunnerOutput(
req_ids=[request.request_id for request in requests],
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([0]), np.array([0])]
- + [np.array([]) for _ in range(len(requests) - 2)],
+ sampled_token_ids=[[0], [0]] + [[] for _ in range(len(requests) - 2)],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -349,8 +347,8 @@ def test_stop_via_update_from_output():
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
sampled_token_ids=[
- np.array([EOS_TOKEN_ID]),
- np.array([10, 11]),
+ [EOS_TOKEN_ID],
+ [10, 11],
], # First request hits EOS, second continues
logprobs=None,
prompt_logprobs_dict={},
@@ -394,10 +392,7 @@ def test_stop_via_update_from_output():
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
- sampled_token_ids=[
- np.array([10, 42, 12]),
- np.array([13, 14]),
- ], # First request hits stop token
+ sampled_token_ids=[[10, 42, 12], [13, 14]], # First request hits stop token
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -441,10 +436,7 @@ def test_stop_via_update_from_output():
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
- sampled_token_ids=[
- np.array([10, 11, 12]),
- np.array([13]),
- ], # First request exceeds max_tokens
+ sampled_token_ids=[[10, 11, 12], [13]], # First request exceeds max_tokens
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -483,7 +475,7 @@ def test_stop_via_update_from_output():
model_output = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
- sampled_token_ids=[np.array([EOS_TOKEN_ID, 10, 11])],
+ sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -590,12 +582,12 @@ def test_check_stop_min_tokens():
@pytest.mark.parametrize(
"enable_prefix_caching, prompt_logprobs",
[
- (None, None),
+ (False, None),
(True, 5),
],
)
def test_schedule_concurrent_batches(
- enable_prefix_caching: bool | None, prompt_logprobs: int | None
+ enable_prefix_caching: bool, prompt_logprobs: int | None
):
scheduler = create_scheduler(
max_num_batched_tokens=1024,
@@ -624,7 +616,7 @@ def test_schedule_concurrent_batches(
model_runner_output = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -641,7 +633,7 @@ def test_schedule_concurrent_batches(
model_runner_output = ModelRunnerOutput(
req_ids=[requests[1].request_id],
req_id_to_index={requests[1].request_id: 0},
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -678,7 +670,7 @@ def test_preempt_during_execution():
model_runner_output0 = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -695,7 +687,7 @@ def test_preempt_during_execution():
model_runner_output1 = ModelRunnerOutput(
req_ids=[requests[1].request_id],
req_id_to_index={requests[1].request_id: 0},
- sampled_token_ids=[np.array([42])],
+ sampled_token_ids=[[42]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -712,18 +704,14 @@ def test_preempt_during_execution():
@pytest.mark.parametrize(
"spec_tokens,output_tokens,expected",
[
- ([[1, 2, 3]], [np.array([1, 2, 3, 4])], (1, 3, 3, [1, 1, 1])), # perfect match
- ([[1, 2, 3]], [np.array([1, 5])], (1, 3, 1, [1, 0, 0])), # early mismatch
- (
- [[1, 2], [3]],
- [np.array([1, 2, 5]), np.array([3, 4])],
- (2, 3, 3, [2, 1]),
- ), # multiple sequences
- ([[1]], [np.array([1, 2])], (1, 1, 1, [1])), # single token sequence
- ([[]], [np.array([5])], (0, 0, 0, [0])), # empty sequence
+ ([[1, 2, 3]], [[1, 2, 3, 4]], (1, 3, 3, [1, 1, 1])), # perfect match
+ ([[1, 2, 3]], [[1, 5]], (1, 3, 1, [1, 0, 0])), # early mismatch
+ ([[1, 2], [3]], [[1, 2, 5], [3, 4]], (2, 3, 3, [2, 1])), # multiple sequences
+ ([[1]], [[1, 2]], (1, 1, 1, [1])), # single token sequence
+ ([[]], [[5]], (0, 0, 0, [0])), # empty sequence
(
[[1, 2, 3], [4, 5, 6]],
- [np.array([1, 2, 7]), np.array([4, 8])],
+ [[1, 2, 7], [4, 8]],
(2, 6, 3, [2, 1, 0]),
), # multiple mismatches
],
@@ -757,7 +745,7 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
model_runner_output = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([0]) for _ in range(len(requests))],
+ sampled_token_ids=[[0] for _ in range(len(requests))],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -984,7 +972,7 @@ def test_kv_connector_basic(is_async: bool):
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([1000])] * len(req_ids),
+ sampled_token_ids=[[1000]] * len(req_ids),
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1037,7 +1025,7 @@ def test_kv_connector_basic(is_async: bool):
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([1000])] * len(req_ids),
+ sampled_token_ids=[[1000]] * len(req_ids),
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1069,7 +1057,8 @@ def test_kv_connector_basic(is_async: bool):
)
-def test_external_prefix_cache_metrics():
+@pytest.mark.parametrize("is_async", [False, True])
+def test_external_prefix_cache_metrics(is_async: bool):
"""
Verify connector prefix cache metrics are updated
correctly when the scheduler processes requests with KV connector hits.
@@ -1079,7 +1068,9 @@ def test_external_prefix_cache_metrics():
NUM_MATCHED_NEW_TOKENS = 4
scheduler = create_scheduler(
enable_prefix_caching=False,
- use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),
+ use_kv_connector=mock_kv(
+ matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async
+ ),
)
# --- Prepare simple requests ---
@@ -1091,16 +1082,22 @@ def test_external_prefix_cache_metrics():
num_tokens=NUM_TOKENS,
max_tokens=MAX_TOKENS,
)
+ req_ids = []
+ req_to_index = {}
+ for i, request in enumerate(requests):
+ scheduler.add_request(request)
+ req_ids.append(request.request_id)
+ req_to_index[request.request_id] = i
- for req in requests:
- scheduler.add_request(req)
+ if is_async:
+ _step_until_kv_transfer_finished(scheduler, req_ids)
# --- Trigger scheduling and simulate model output ---
output = scheduler.schedule()
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
req_ids=[r.request_id for r in requests],
req_id_to_index={r.request_id: i for i, r in enumerate(requests)},
- sampled_token_ids=[np.array([1000])] * NUM_REQUESTS,
+ sampled_token_ids=[[1000]] * NUM_REQUESTS,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1166,7 +1163,7 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([1000])] * len(req_ids),
+ sampled_token_ids=[[1000]] * len(req_ids),
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1251,7 +1248,7 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([1000])] * len(req_ids),
+ sampled_token_ids=[[1000]] * len(req_ids),
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1344,7 +1341,7 @@ def make_output(scheduler: Scheduler):
return ModelRunnerOutput(
req_ids=[req.request_id for req in scheduler.running],
req_id_to_index={req.request_id: i for i, req in enumerate(scheduler.running)},
- sampled_token_ids=[np.array([1000])] * len(scheduler.running),
+ sampled_token_ids=[[1000]] * len(scheduler.running),
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1428,7 +1425,7 @@ def create_scheduler_with_priority(
model: str = "facebook/opt-125m",
max_num_seqs: int = 16,
max_num_batched_tokens: int = 8192,
- enable_prefix_caching: bool | None = None,
+ enable_prefix_caching: bool = False,
long_prefill_token_threshold: int = 0,
disable_chunked_mm_input: bool = False,
use_kv_connector: bool = False,
@@ -1447,7 +1444,7 @@ def create_scheduler_with_priority(
max_num_batch_tokens: max num tokens to batch
enable_prefix_caching: optionally force APC config
(True/False) or use default
- (None)
+ (False)
Returns:
{class}`Scheduler` instance with priority scheduling
@@ -1470,17 +1467,12 @@ def create_scheduler_with_priority(
seed=42,
)
# Cache config, optionally force APC
- kwargs_cache = (
- {}
- if enable_prefix_caching is None
- else {"enable_prefix_caching": enable_prefix_caching}
- )
cache_config = CacheConfig(
block_size=block_size,
gpu_memory_utilization=0.9,
swap_space=0,
cache_dtype="auto",
- **kwargs_cache,
+ enable_prefix_caching=enable_prefix_caching,
)
kv_transfer_config = (
KVTransferConfig(
@@ -1761,7 +1753,7 @@ def test_priority_scheduling_preemption():
req_id_to_index={
req.request_id: i for i, req in enumerate(low_priority_requests)
},
- sampled_token_ids=[np.array([100]) for _ in low_priority_requests],
+ sampled_token_ids=[[100] for _ in low_priority_requests],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -1830,7 +1822,7 @@ def test_priority_scheduling_no_preemption_when_space_available():
req_id_to_index={
req.request_id: i for i, req in enumerate(low_priority_requests)
},
- sampled_token_ids=[np.array([100]) for _ in low_priority_requests],
+ sampled_token_ids=[[100] for _ in low_priority_requests],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -2076,7 +2068,7 @@ def test_priority_scheduling_heap_property():
model_output = ModelRunnerOutput(
req_ids=[req.req_id],
req_id_to_index={req.req_id: 0},
- sampled_token_ids=[np.array([100])],
+ sampled_token_ids=[[100]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -2162,7 +2154,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
model_output = ModelRunnerOutput(
req_ids=[request_low.request_id],
req_id_to_index={request_low.request_id: 0},
- sampled_token_ids=[np.array([100])],
+ sampled_token_ids=[[100]],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -2193,7 +2185,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
- sampled_token_ids=[np.array([100]) for _ in requests],
+ sampled_token_ids=[[100] for _ in requests],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -2219,7 +2211,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
- sampled_token_ids=[np.array([]), np.array([100])],
+ sampled_token_ids=[[], [100]],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -2636,7 +2628,7 @@ def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):
model_output = ModelRunnerOutput(
req_ids=[request1.request_id],
req_id_to_index={request1.request_id: 0},
- sampled_token_ids=[np.array([100])],
+ sampled_token_ids=[[100]],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -2842,7 +2834,7 @@ def test_ec_connector_unable_to_allocate(use_kv_connector):
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
- sampled_token_ids=[np.array([1000])] * len(req_ids),
+ sampled_token_ids=[[1000]] * len(req_ids),
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
@@ -2955,7 +2947,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
model_output = ModelRunnerOutput(
req_ids=[request_low.request_id],
req_id_to_index={request_low.request_id: 0},
- sampled_token_ids=[np.array([100])],
+ sampled_token_ids=[[100]],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -3006,7 +2998,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
- sampled_token_ids=[np.array([100]) for _ in requests],
+ sampled_token_ids=[[100] for _ in requests],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -3041,7 +3033,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
- sampled_token_ids=[np.array([100]), np.array([100, 200])],
+ sampled_token_ids=[[100], [100, 200]],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
@@ -3227,7 +3219,7 @@ def test_ec_connector_allocate_encoder_tokens_with_external_load(use_kv_connecto
model_output = ModelRunnerOutput(
req_ids=[request1.request_id, request2.request_id],
req_id_to_index={request1.request_id: 0, request2.request_id: 1},
- sampled_token_ids=[np.array([100]), np.array([121])],
+ sampled_token_ids=[[100], [121]],
# spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py
index 65511c17473b2..6830f68736453 100644
--- a/tests/v1/core/utils.py
+++ b/tests/v1/core/utils.py
@@ -42,7 +42,7 @@ def create_scheduler(
model: str = "facebook/opt-125m",
max_num_seqs: int = 16,
max_num_batched_tokens: int = 8192,
- enable_prefix_caching: bool | None = None,
+ enable_prefix_caching: bool = False,
long_prefill_token_threshold: int = 0,
disable_chunked_mm_input: bool = False,
use_kv_connector: None | bool | MockKVConfig = None,
@@ -63,7 +63,7 @@ def create_scheduler(
max_num_batch_tokens: max num tokens to batch
enable_prefix_caching: optionally force APC config
(True/False) or use default
- (None)
+ (False)
Returns:
{class}`Scheduler` instance
@@ -87,17 +87,12 @@ def create_scheduler(
skip_tokenizer_init=skip_tokenizer_init,
)
# Cache config, optionally force APC
- kwargs_cache = (
- {}
- if enable_prefix_caching is None
- else {"enable_prefix_caching": enable_prefix_caching}
- )
cache_config = CacheConfig(
block_size=block_size,
gpu_memory_utilization=0.9,
swap_space=0,
cache_dtype="auto",
- **kwargs_cache,
+ enable_prefix_caching=enable_prefix_caching,
)
kv_transfer_config = None
if isinstance(use_kv_connector, MockKVConfig):
diff --git a/tests/v1/determinism/conftest.py b/tests/v1/determinism/conftest.py
index 3c2136e005849..bde02bbd0d5c6 100644
--- a/tests/v1/determinism/conftest.py
+++ b/tests/v1/determinism/conftest.py
@@ -1,11 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
import pytest
+import vllm.model_executor.layers.batch_invariant as batch_invariant
+
@pytest.fixture(autouse=True)
def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):
"""Automatically enable batch invariant kernel overrides for all tests."""
+ monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
- yield
diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py
index f018ee551dbfe..b9e2daafb8705 100644
--- a/tests/v1/determinism/test_batch_invariance.py
+++ b/tests/v1/determinism/test_batch_invariance.py
@@ -6,8 +6,15 @@ import random
import pytest
import torch
-from utils import _extract_step_logprobs, _random_prompt, skip_unsupported
+from utils import (
+ BACKENDS,
+ _extract_step_logprobs,
+ _random_prompt,
+ resolve_model_name,
+ skip_unsupported,
+)
+import vllm.model_executor.layers.batch_invariant as batch_invariant
from vllm import LLM, SamplingParams
@@ -15,7 +22,7 @@ from vllm import LLM, SamplingParams
@pytest.mark.timeout(1000)
@pytest.mark.parametrize(
"backend",
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
+ BACKENDS,
)
def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
backend, monkeypatch: pytest.MonkeyPatch
@@ -47,7 +54,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
# Allow overrides from environment (useful for CI tuning)
# "facebook/opt-125m" is too small, doesn't reliably test determinism
- model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
+ model = resolve_model_name(backend)
num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))
max_batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "128"))
min_random_prompt = int(os.getenv("VLLM_MIN_PROMPT", "1024"))
@@ -150,7 +157,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
@skip_unsupported
@pytest.mark.parametrize(
"backend",
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
+ BACKENDS,
)
@pytest.mark.forked
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
@@ -160,7 +167,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
random.seed(seed)
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
+ model_name = resolve_model_name(backend)
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
# For batch invariance, disable custom all-reduce to ensure deterministic
@@ -183,6 +190,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
max_num_seqs=32,
max_model_len=8192,
dtype="bfloat16", # not everything is supported
+ gpu_memory_utilization=0.9,
)
# Use more realistic prompts for better token generation
@@ -369,7 +377,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
@skip_unsupported
@pytest.mark.parametrize(
"backend",
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
+ BACKENDS,
)
def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
"""
@@ -377,7 +385,7 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
Useful for quick smoke testing and debugging.
"""
monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
- model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
+ model = resolve_model_name(backend)
llm = LLM(
model=model,
@@ -419,7 +427,7 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
@skip_unsupported
@pytest.mark.parametrize(
"backend",
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
+ BACKENDS,
)
@pytest.mark.forked
def test_logprobs_without_batch_invariance_should_fail(
@@ -438,10 +446,10 @@ def test_logprobs_without_batch_invariance_should_fail(
# CRITICAL: Disable batch invariance for this test
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
-
+ monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", False)
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
random.seed(seed)
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
+ model_name = resolve_model_name(backend)
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
print(f"\n{'=' * 80}")
@@ -659,7 +667,7 @@ def test_decode_logprobs_match_prefill_logprobs(
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
random.seed(seed)
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
+ model_name = resolve_model_name(backend)
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
from vllm.model_executor.layers.batch_invariant import (
diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py
index 23f47863dd23f..d74b435797f8f 100644
--- a/tests/v1/determinism/test_online_batch_invariance.py
+++ b/tests/v1/determinism/test_online_batch_invariance.py
@@ -16,7 +16,8 @@ import sys
from typing import Any
import openai
-from utils import _random_prompt, skip_unsupported
+import pytest
+from utils import BACKENDS, _random_prompt, resolve_model_name, skip_unsupported
from tests.utils import RemoteOpenAIServer
@@ -133,9 +134,14 @@ def _compare_bs1_vs_bsn_single_process(
@skip_unsupported
-def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN():
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
+ backend: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
+ # Override backend for this test (and the RemoteOpenAIServer child process).
+ monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
+ model_name = resolve_model_name(backend)
prompts_all = [_random_prompt(10, 50) for _ in range(32)]
sp_kwargs: dict[str, Any] = {
diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py
index 5141837faea04..ecbb6a1126933 100644
--- a/tests/v1/determinism/utils.py
+++ b/tests/v1/determinism/utils.py
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
import random
import pytest
import torch
+from vllm.attention.utils.fa_utils import flash_attn_supports_mla
from vllm.platforms import current_platform
skip_unsupported = pytest.mark.skipif(
@@ -12,6 +14,25 @@ skip_unsupported = pytest.mark.skipif(
reason="Requires CUDA and >= Hopper (SM90)",
)
+BACKENDS: list[str] = [
+ "FLASH_ATTN",
+ "FLASHINFER",
+]
+
+if flash_attn_supports_mla():
+ BACKENDS.append("FLASH_ATTN_MLA")
+
+DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
+MLA_MODEL = "deepseek-ai/DeepSeek-V2-Lite-Chat"
+
+
+def resolve_model_name(backend: str) -> str:
+ """Resolve the model name for the given backend."""
+ model = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
+ if backend.endswith("MLA") and model == DEFAULT_MODEL:
+ return MLA_MODEL
+ return model
+
def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
# Generate more realistic prompts that will actually produce varied tokens
diff --git a/tests/v1/e2e/test_async_scheduling.py b/tests/v1/e2e/test_async_scheduling.py
index f732b05f09f9d..00d93e1ba0b53 100644
--- a/tests/v1/e2e/test_async_scheduling.py
+++ b/tests/v1/e2e/test_async_scheduling.py
@@ -84,6 +84,7 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):
"num_speculative_tokens": 2,
"model": "nm-testing/Llama3_2_1B_speculator.eagle3",
}
+ # Set small draft model len to force doesn't-fit-in-drafter case.
spec_config_short = spec_config | {"max_model_len": 50}
# test_preemption, executor, async_scheduling,
@@ -174,13 +175,14 @@ def run_tests(
):
if "spec_mml=None" in test_config:
assert (
- pytest.approx(test_acceptance_rate, rel=5e-2)
- == base_acceptance_rate
+ test_acceptance_rate > base_acceptance_rate
+ or test_acceptance_rate
+ == pytest.approx(base_acceptance_rate, rel=5e-2)
)
else:
# Currently the reported acceptance rate is expected to be
# lower when we sometimes skip drafting altogether.
- assert test_acceptance_rate > 0.05
+ assert test_acceptance_rate > 0.1
print(
f"PASSED: config=[{test_config}], params={params}"
f" accept_rate={test_acceptance_rate}"
diff --git a/tests/v1/e2e/test_lora_with_spec_decode.py b/tests/v1/e2e/test_lora_with_spec_decode.py
index 14532f2795443..8c9ab58c3c0ab 100644
--- a/tests/v1/e2e/test_lora_with_spec_decode.py
+++ b/tests/v1/e2e/test_lora_with_spec_decode.py
@@ -61,8 +61,6 @@ def test_batch_inference_correctness(
model_setup: (method, model_name, spec_model_name, lora_path, tp_size)
"""
with monkeypatch.context() as m:
- m.setenv("VLLM_USE_V1", "1")
-
# Disable randomness
m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
torch.manual_seed(SEED)
diff --git a/tests/v1/ec_connector/integration/README.md b/tests/v1/ec_connector/integration/README.md
index 30426e055ade8..2dbcb307fda32 100644
--- a/tests/v1/ec_connector/integration/README.md
+++ b/tests/v1/ec_connector/integration/README.md
@@ -113,7 +113,7 @@ Quick sanity check:
- Outputs differ between baseline and disagg
- Server startup fails
-- Encoder cache not found (should fallback to local execution)
+- Encoder cache not found (should fall back to local execution)
- Proxy routing errors
## Notes
diff --git a/tests/v1/entrypoints/llm/test_struct_output_generate.py b/tests/v1/entrypoints/llm/test_struct_output_generate.py
index a7d769c8542a9..d1b037b7956cf 100644
--- a/tests/v1/entrypoints/llm/test_struct_output_generate.py
+++ b/tests/v1/entrypoints/llm/test_struct_output_generate.py
@@ -46,17 +46,45 @@ EAGLE_SPEC_CONFIG = {
PARAMS_MODELS_BACKENDS_TOKENIZER_MODE = [
("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "auto", None),
- ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto", None),
- ("mistralai/Ministral-8B-Instruct-2410", "lm-format-enforcer", "auto", None),
+ # FIXME: Since "auto" will use Mistral tokenizer and these backends do not support
+ # it, we skip these tests for now.
+ # ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto", None),
+ # ("mistralai/Ministral-8B-Instruct-2410", "lm-format-enforcer", "auto", None),
+ ("mistralai/Ministral-8B-Instruct-2410", "guidance", "hf", None),
+ pytest.param(
+ "mistralai/Ministral-8B-Instruct-2410",
+ "lm-format-enforcer",
+ "hf",
+ None,
+ marks=pytest.mark.skip(
+ reason=(
+ "Flaky: lm-format-enforcer intermittently returns"
+ "incomplete JSON."
+ "See https://github.com/noamgat/lm-format-enforcer/issues/169"
+ )
+ ),
+ ),
("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "mistral", None),
("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", None),
- ("Qwen/Qwen2.5-1.5B-Instruct", "lm-format-enforcer", "auto", None),
+ pytest.param(
+ "Qwen/Qwen2.5-1.5B-Instruct",
+ "lm-format-enforcer",
+ "auto",
+ None,
+ marks=pytest.mark.skip(
+ reason=(
+ "Flaky: lm-format-enforcer intermittently returns"
+ "incomplete JSON."
+ "See https://github.com/noamgat/lm-format-enforcer/issues/169"
+ )
+ ),
+ ),
# FIXME: This tests are flaky on CI thus disabled. Tracking in Issue #24402
# ("mistralai/Ministral-8B-Instruct-2410", "outlines", "auto", None),
# ("mistralai/Ministral-8B-Instruct-2410", "outlines", "mistral", None),
# ("Qwen/Qwen2.5-1.5B-Instruct", "guidance", "auto"),
("mistralai/Ministral-8B-Instruct-2410", "outlines", "auto", NGRAM_SPEC_CONFIG),
- ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto", NGRAM_SPEC_CONFIG),
+ ("mistralai/Ministral-8B-Instruct-2410", "guidance", "hf", NGRAM_SPEC_CONFIG),
("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", NGRAM_SPEC_CONFIG),
("meta-llama/Meta-Llama-3.1-8B-Instruct", "xgrammar", "auto", EAGLE_SPEC_CONFIG),
]
@@ -97,7 +125,6 @@ def test_guided_decoding_deprecated():
assert sp1.structured_outputs == guided_decoding
-@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize(
"model_name, backend, tokenizer_mode, speculative_config",
PARAMS_MODELS_BACKENDS_TOKENIZER_MODE,
@@ -128,6 +155,8 @@ def test_structured_output(
),
seed=120,
tokenizer_mode=tokenizer_mode,
+ load_format="auto" if not model_name.startswith("mistralai/") else "hf",
+ config_format="auto" if not model_name.startswith("mistralai/") else "hf",
speculative_config=speculative_config,
)
@@ -602,7 +631,6 @@ Make the response as short as possible.
)
-@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize(
"model_name, backend, tokenizer_mode, reasoning_parser, speculative_config", # noqa: E501
[
@@ -687,7 +715,6 @@ def test_structured_output_with_reasoning_matrices(
jsonschema.validate(instance=output_json, schema=reasoning_schema)
-@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("model_name, tokenizer_mode", PARAMS_MODELS_TOKENIZER_MODE)
def test_structured_output_auto_mode(
unsupported_json_schema: dict[str, Any],
@@ -699,6 +726,8 @@ def test_structured_output_auto_mode(
max_model_len=1024,
structured_outputs_config=dict(backend="auto"),
tokenizer_mode=tokenizer_mode,
+ load_format="auto",
+ config_format="auto",
)
sampling_params = SamplingParams(
@@ -734,7 +763,6 @@ def test_structured_output_auto_mode(
assert isinstance(parsed_json, dict)
-@pytest.mark.skip_global_cleanup
def test_guidance_no_additional_properties():
llm = LLM(
model="Qwen/Qwen2.5-1.5B-Instruct",
diff --git a/tests/v1/entrypoints/openai/serving_responses/test_image.py b/tests/v1/entrypoints/openai/serving_responses/test_image.py
index 980d83b787e7a..be5693bbf2736 100644
--- a/tests/v1/entrypoints/openai/serving_responses/test_image.py
+++ b/tests/v1/entrypoints/openai/serving_responses/test_image.py
@@ -15,10 +15,10 @@ MODEL_NAME = "Qwen/Qwen2.5-VL-3B-Instruct"
MAXIMUM_IMAGES = 2
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
- "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
- "Grayscale_8bits_palette_sample_image.png", # "https://upload.wikimedia.org/wikipedia/commons/f/fa/Grayscale_8bits_palette_sample_image.png",
- "1280px-Venn_diagram_rgb.svg.png", # "https://upload.wikimedia.org/wikipedia/commons/thumb/9/91/Venn_diagram_rgb.svg/1280px-Venn_diagram_rgb.svg.png",
- "RGBA_comp.png", # "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
+ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+ "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
+ "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",
]
diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
index ebc8575e5b390..453ccc81eb14a 100755
--- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
@@ -49,13 +49,13 @@ NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1
PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
-PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-16}
-DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-16}
+PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
+DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
# Find the git repository root directory
GIT_ROOT=$(git rev-parse --show-toplevel)
-SMI_BIN=$(which nvidia-smi || which rocm-smi)
+SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
# Trap the SIGINT signal (triggered by Ctrl+C)
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
@@ -91,8 +91,13 @@ get_model_args() {
get_num_gpus() {
if [[ "$SMI_BIN" == *"nvidia"* ]]; then
echo "$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)"
- else
+ elif [[ "$SMI_BIN" == *"rocm"* ]]; then
echo "$($SMI_BIN -l | grep GPU | wc -l)"
+ else
+ # works for non-cuda platforms,
+ # assuming at least 1 device and
+ # let system to decide which card to use
+ echo "1"
fi
}
diff --git a/tests/v1/kv_connector/unit/test_lmcache_integration.py b/tests/v1/kv_connector/unit/test_lmcache_integration.py
index 11507d7cd4e7b..33418edc325af 100644
--- a/tests/v1/kv_connector/unit/test_lmcache_integration.py
+++ b/tests/v1/kv_connector/unit/test_lmcache_integration.py
@@ -9,6 +9,12 @@
# Assumption vs. Correctness Tests:
# these unit tests do *not* test correctness of LMCache-side or vLLM-side logic
# it is to ensure that assumptions LMCache makes about vLLM's interface are stable
+
+import pytest
+
+from vllm.platforms import current_platform
+
+
def assumes(obj, attr, is_callable=False, is_instance_of=None):
import inspect
from dataclasses import is_dataclass
@@ -48,6 +54,9 @@ def assumes(obj, attr, is_callable=False, is_instance_of=None):
assert isinstance(attr_value, is_instance_of), assumption_msg
+@pytest.mark.skipif(
+ current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
+)
def test_multimodal_interface():
# protect against interface changes
from vllm.multimodal.inputs import PlaceholderRange
@@ -72,6 +81,9 @@ def test_multimodal_interface():
assert token_ids.tolist() == [0, 0, 0, 0, 4, 4369, 4369, 4369, 4369, 9]
+@pytest.mark.skipif(
+ current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
+)
def test_config_interface():
# protect against interface changes
from vllm.config import VllmConfig
@@ -146,6 +158,9 @@ def test_config_interface():
)
+@pytest.mark.skipif(
+ current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
+)
def test_request_interface():
# protect against interface changes
from types import NoneType
diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py
index 1c1ac915c758e..ffa7d884d2762 100644
--- a/tests/v1/kv_connector/unit/test_multi_connector.py
+++ b/tests/v1/kv_connector/unit/test_multi_connector.py
@@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
NixlKVConnectorStats,
)
+from vllm.platforms import current_platform
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
@@ -69,6 +70,13 @@ def _compare_directories(dir1: Path, dir2: Path) -> bool:
return True
+@pytest.mark.skipif(
+ current_platform.is_rocm(),
+ reason=(
+ "hipErrorLaunchFailure when running this test, see issue:"
+ "https://github.com/ROCm/pytorch/issues/2822"
+ ),
+)
def test_multi_shared_storage_connector_consistency():
"""
Tests that MultiConnector with two SharedStorageConnectors saves
diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py
index b264e5108c16d..b7d7a10057b8b 100644
--- a/tests/v1/kv_connector/unit/test_nixl_connector.py
+++ b/tests/v1/kv_connector/unit/test_nixl_connector.py
@@ -11,7 +11,6 @@ import uuid
from collections import defaultdict
from unittest.mock import patch
-import numpy as np
import pytest
import ray
import torch
@@ -827,7 +826,7 @@ def test_kv_connector_stats_aggregation():
output = ModelRunnerOutput(
req_ids=[f"req_{i}"],
req_id_to_index={f"req_{i}": 0},
- sampled_token_ids=[np.array([123])], # dummy token
+ sampled_token_ids=[[123]], # dummy token
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[None],
@@ -908,7 +907,7 @@ def test_multi_kv_connector_stats_aggregation():
output = ModelRunnerOutput(
req_ids=[f"req_{i}"],
req_id_to_index={f"req_{i}": 0},
- sampled_token_ids=[np.array([123])],
+ sampled_token_ids=[[123]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[None],
@@ -966,7 +965,7 @@ def test_scheduler_kv_connector_stats_aggregation():
model_output = ModelRunnerOutput(
req_ids=["req_0"],
req_id_to_index={"req_0": 0},
- sampled_token_ids=[np.array([123])],
+ sampled_token_ids=[[123]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[None],
diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py
index 23b6c4802d106..69565f584ab89 100644
--- a/tests/v1/kv_connector/unit/test_offloading_connector.py
+++ b/tests/v1/kv_connector/unit/test_offloading_connector.py
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
)
from vllm.forward_context import ForwardContext
from vllm.utils.hashing import sha256
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
from vllm.v1.core.kv_cache_utils import (
BlockHash,
get_request_block_hasher,
@@ -92,7 +93,7 @@ class MockOffloadingSpec(OffloadingSpec):
return self.manager
def get_handlers(
- self, _
+ self, _, __
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
@@ -138,7 +139,10 @@ class RequestRunner:
self.worker_connector = OffloadingConnector(vllm_config, KVConnectorRole.WORKER)
# register worker kv_caches to enable OffloadingWorker creations
- self.worker_connector.register_kv_caches(kv_caches={"a": torch.empty(0)})
+ self.worker_connector.register_cross_layers_kv_cache(
+ kv_cache=torch.empty(0),
+ attn_backend=FlashAttentionBackend,
+ )
# extract connector of scheduler
scheduler_connector = self.scheduler.connector
diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py
index c248104d5b5ea..f35f91bb3adf8 100644
--- a/tests/v1/kv_connector/unit/utils.py
+++ b/tests/v1/kv_connector/unit/utils.py
@@ -7,7 +7,6 @@ from dataclasses import dataclass
from itertools import chain, count
from typing import Any
-import numpy as np
import torch
from vllm import SamplingParams
@@ -229,7 +228,7 @@ def create_model_runner_output(
# Make sampled tokens.
sampled_token = EOS_TOKEN_ID if use_eos else token_id
- sampled_token_ids = [np.array([sampled_token]) for _ in req_ids]
+ sampled_token_ids = [[sampled_token] for _ in req_ids]
kv_connector_output = (
None
diff --git a/tests/v1/kv_offload/test_cpu_gpu.py b/tests/v1/kv_offload/test_cpu_gpu.py
index 0d4fa344d298c..a248104e16d2d 100644
--- a/tests/v1/kv_offload/test_cpu_gpu.py
+++ b/tests/v1/kv_offload/test_cpu_gpu.py
@@ -103,8 +103,8 @@ def test_transfer(
for i in range(gpu_blocks_per_cpu_block):
cpu_blocks_in_gpu_block_size.append(i + base_block_id)
- # maybe skip a GPU block to test writing to the middle of a CPU block
- if gpu_to_cpu:
+ # maybe skip a GPU block to test reading from the middle of a CPU block
+ if not gpu_to_cpu:
gpu_blocks = gpu_blocks[gpu_blocks_per_cpu_block - 1 :]
cpu_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size[
gpu_blocks_per_cpu_block - 1 :
diff --git a/tests/v1/kv_offload/test_cpu_offloading.py b/tests/v1/kv_offload/test_cpu_offloading.py
index b654ea4298dbb..406d4c0b4c1fd 100644
--- a/tests/v1/kv_offload/test_cpu_offloading.py
+++ b/tests/v1/kv_offload/test_cpu_offloading.py
@@ -12,8 +12,14 @@ from tqdm import tqdm
from vllm import LLM, SamplingParams, TokensPrompt
from vllm.config import KVEventsConfig, KVTransferConfig
from vllm.distributed.kv_events import BlockStored, KVEventBatch
+from vllm.platforms import current_platform
+from vllm.utils.system_utils import set_env_var
-CPU_BLOCK_SIZES = [16, 48]
+CPU_BLOCK_SIZES = [48]
+ATTN_BACKENDS = ["FLASH_ATTN"]
+
+if current_platform.is_cuda():
+ ATTN_BACKENDS.append("FLASHINFER")
class MockSubscriber:
@@ -63,8 +69,88 @@ class MockSubscriber:
self.sub.close()
+def _latency_test(llm: LLM, subscriber: MockSubscriber):
+ sampling_params = SamplingParams(max_tokens=1)
+
+ num_times_cpu_better_than_cold = 0
+ num_tests = 10
+ total_cold_time = 0.0
+ total_gpu_hit_time = 0.0
+ total_cpu_hit_time = 0.0
+ prompt_token_ids = [0] * 10001
+ for i in tqdm(range(num_tests), desc="Running tests"):
+ prompt_token_ids[0] = i
+ prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
+
+ # run generation - this should trigger saving KV cache
+ start_time = time.time()
+ llm.generate(prompts, sampling_params, use_tqdm=False)
+ cold_time = time.time() - start_time
+ total_cold_time += cold_time
+
+ # run generation again - should hit the GPU prefix cache
+ start_time = time.time()
+ llm.generate(prompts, sampling_params, use_tqdm=False)
+ gpu_hit_time = time.time() - start_time
+ total_gpu_hit_time += gpu_hit_time
+
+ # reset prefix cache to avoid GPU hit.
+ llm.reset_prefix_cache()
+
+ assert subscriber.get_new_cpu_stored_events()
+
+ # run generation again - this should trigger loading from CPU
+ start_time = time.time()
+ llm.generate(prompts, sampling_params, use_tqdm=False)
+ cpu_hit_time = time.time() - start_time
+ total_cpu_hit_time += cpu_hit_time
+
+ if cpu_hit_time < cold_time:
+ num_times_cpu_better_than_cold += 1
+
+ print("Average times:")
+ print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms")
+ print(f" GPU hit: {total_gpu_hit_time * 1000 / num_tests:.2f}ms")
+ print(f" CPU hit: {total_cpu_hit_time * 1000 / num_tests:.2f}ms")
+
+ assert num_times_cpu_better_than_cold >= 0.8 * num_tests
+
+
+def _accuracy_test(llm: LLM, subscriber: MockSubscriber):
+ sampling_params = SamplingParams(max_tokens=1)
+ cpu_block_size = (
+ llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config[
+ "block_size"
+ ]
+ )
+
+ subscriber.get_new_cpu_stored_events()
+
+ # prepend prompt to be cpu block aligned
+ prompt = "Let's count to 10. One, two, three, four,"
+ while (
+ len(llm.generate(prompt, use_tqdm=False)[0].prompt_token_ids) % cpu_block_size
+ != 0
+ ):
+ prompt = ". " + prompt
+
+ assert subscriber.get_new_cpu_stored_events()
+
+ test_count = 100
+ success_count = 0
+ for i in range(test_count):
+ if (
+ llm.generate(prompt, sampling_params, use_tqdm=False)[0].outputs[0].text
+ == " five"
+ ):
+ success_count += 1
+
+ assert success_count >= 0.5 * test_count
+
+
@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
-def test_cpu_offloading(cpu_block_size: int) -> None:
+@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
+def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
"""
Tests OffloadingConnector with CPUOffloadingSpec.
"""
@@ -92,61 +178,20 @@ def test_cpu_offloading(cpu_block_size: int) -> None:
topic="test",
)
- llm = LLM(
- model="meta-llama/Llama-3.2-1B-Instruct",
- gpu_memory_utilization=0.5,
- kv_events_config=kv_events_config,
- kv_transfer_config=kv_transfer_config,
- )
-
- sampling_params = SamplingParams(temperature=0, max_tokens=1)
+ with set_env_var("VLLM_ATTENTION_BACKEND", attn_backend):
+ llm = LLM(
+ model="meta-llama/Llama-3.2-1B-Instruct",
+ gpu_memory_utilization=0.5,
+ kv_events_config=kv_events_config,
+ kv_transfer_config=kv_transfer_config,
+ )
events_endpoint = events_endpoint.replace("*", "127.0.0.1")
subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
try:
- num_times_cpu_better_than_cold = 0
- num_tests = 10
- total_cold_time = 0.0
- total_gpu_hit_time = 0.0
- total_cpu_hit_time = 0.0
- prompt_token_ids = [0] * 10001
- for i in tqdm(range(num_tests), desc="Running tests"):
- prompt_token_ids[0] = i
- prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
-
- # run generation - this should trigger saving KV cache
- start_time = time.time()
- llm.generate(prompts, sampling_params, use_tqdm=False)
- cold_time = time.time() - start_time
- total_cold_time += cold_time
-
- # run generation again - should hit the GPU prefix cache
- start_time = time.time()
- llm.generate(prompts, sampling_params, use_tqdm=False)
- gpu_hit_time = time.time() - start_time
- total_gpu_hit_time += gpu_hit_time
-
- # reset prefix cache to avoid GPU hit.
- llm.reset_prefix_cache()
-
- assert subscriber.get_new_cpu_stored_events()
-
- # run generation again - this should trigger loading from CPU
- start_time = time.time()
- llm.generate(prompts, sampling_params, use_tqdm=False)
- cpu_hit_time = time.time() - start_time
- total_cpu_hit_time += cpu_hit_time
-
- if cpu_hit_time < cold_time:
- num_times_cpu_better_than_cold += 1
-
- print("Average times:")
- print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms")
- print(f" GPU hit: {total_gpu_hit_time * 1000 / num_tests:.2f}ms")
- print(f" CPU hit: {total_cpu_hit_time * 1000 / num_tests:.2f}ms")
-
- assert num_times_cpu_better_than_cold >= 0.8 * num_tests
+ _latency_test(llm, subscriber)
+ _accuracy_test(llm, subscriber)
finally:
subscriber.close()
del llm
diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py
index 42584938bc06f..c89c33be80c10 100644
--- a/tests/v1/sample/test_logprobs.py
+++ b/tests/v1/sample/test_logprobs.py
@@ -521,8 +521,8 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):
pytest.param(
(
"eagle",
- "meta-llama/Llama-3.1-8B-Instruct",
- "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
+ "meta-llama/Llama-3.2-1B-Instruct",
+ "nm-testing/Llama3_2_1B_speculator.eagle3",
),
marks=large_gpu_mark(min_gb=32),
),
@@ -541,7 +541,7 @@ def test_spec_decode_logprobs(
"""
from vllm import LLM
- prompt = "Hello world"
+ prompt = "Hello world " * 50
sampling_params = SamplingParams(
temperature=0, logprobs=3, max_tokens=10, ignore_eos=False
)
@@ -582,6 +582,9 @@ def test_spec_decode_logprobs(
seed=42,
logprobs_mode=logprobs_mode,
gpu_memory_utilization=0.4,
+ # Force prefill chunking
+ enable_chunked_prefill=True,
+ max_num_batched_tokens=32,
)
spec_results = spec_llm.generate([prompt], sampling_params)
# Collect logprobs outputs from spec decode LLM.
@@ -597,6 +600,84 @@ def test_spec_decode_logprobs(
# Per-token logprobs are expected to be the same.
assert len(ref_logprobs) == len(spec_logprobs)
for ref_logprob, spec_logprob in zip(ref_logprobs, spec_logprobs):
- assert math.isclose(ref_logprob.logprob, spec_logprob.logprob, abs_tol=1e-3)
+ assert math.isclose(
+ ref_logprob.logprob, spec_logprob.logprob, rel_tol=5e-2, abs_tol=1e-1
+ )
assert ref_logprob.rank == spec_logprob.rank
assert ref_logprob.decoded_token == spec_logprob.decoded_token
+
+
+def test_prompt_logprobs_with_chunking_and_preemption():
+ """Test that prompt logprobs are correctly returned when using
+ both chunked prefill and preemption.
+
+ This test ensures that the num_prompt_logprobs tracking persists
+ across preemptions and prefill chunks.
+ """
+
+ # Create prompts that will trigger chunking and preemption
+ prompts = [
+ "The following numbers of the sequence "
+ + ", ".join(str(i) for i in range(10))
+ + " are:",
+ "In one word, the capital of France is ",
+ ] + [f"Tell me about the number {i}: " for i in range(32)]
+
+ sampling_params = SamplingParams(
+ temperature=0.0,
+ max_tokens=40,
+ min_tokens=20,
+ prompt_logprobs=2, # Request prompt logprobs
+ )
+
+ with VllmRunner(
+ "Qwen/Qwen3-0.6B",
+ max_model_len=512,
+ enable_chunked_prefill=True,
+ max_num_batched_tokens=48, # Force prefill chunking
+ num_gpu_blocks_override=32, # Force preemptions
+ disable_log_stats=False,
+ gpu_memory_utilization=0.25,
+ ) as vllm_model:
+ metrics_before = vllm_model.llm.get_metrics()
+
+ # Generate with prompt logprobs using generate_w_logprobs which
+ # returns (output_ids, output_str, output_logprobs, prompt_logprobs)
+ outputs = vllm_model.generate_w_logprobs(
+ prompts, sampling_params=sampling_params, include_prompt_token_ids=True
+ )
+
+ # Verify that all outputs have prompt logprobs
+ for i, output in enumerate(outputs):
+ _, _, _, prompt_token_ids, prompt_logprobs = output
+ assert prompt_logprobs is not None and len(prompt_logprobs) > 0, (
+ f"Output {i} missing prompt logprobs"
+ )
+ assert len(prompt_logprobs) == len(prompt_token_ids), (
+ "Unexpected number of prompt logprob positions"
+ )
+
+ # Each position should have the requested number of logprobs
+ for pos, logprobs_dict in enumerate(prompt_logprobs):
+ if logprobs_dict is not None: # First token may be None
+ assert (
+ sampling_params.prompt_logprobs
+ <= len(logprobs_dict)
+ <= sampling_params.prompt_logprobs + 1
+ ), (
+ f"Output {i} position {pos} has {len(logprobs_dict)} "
+ f"logprobs, expected {sampling_params.prompt_logprobs}"
+ )
+
+ # Check that we actually had preemptions
+ metrics_after = vllm_model.llm.get_metrics()
+ preemptions_before = next(
+ (m.value for m in metrics_before if m.name == "vllm:num_preemptions"), 0
+ )
+ preemptions_after = next(
+ (m.value for m in metrics_after if m.name == "vllm:num_preemptions"), 0
+ )
+ preemptions = preemptions_after - preemptions_before
+ assert preemptions > 0, "Test did not trigger any preemptions"
+
+ print(f"Test passed with {preemptions} preemptions")
diff --git a/tests/v1/sample/test_sampling_params_e2e.py b/tests/v1/sample/test_sampling_params_e2e.py
index 915b9957031d8..1684252174d3d 100644
--- a/tests/v1/sample/test_sampling_params_e2e.py
+++ b/tests/v1/sample/test_sampling_params_e2e.py
@@ -22,14 +22,6 @@ def test_n_gt_1(llm):
assert len(outputs[0].outputs) == 3
-def test_best_of(llm):
- """Raise a ValueError since best_of is deprecated."""
-
- params = SamplingParams(n=2, best_of=3)
- with pytest.raises(ValueError):
- _ = llm.generate(PROMPT, params)
-
-
def test_penalties(llm):
"""Check that we do not get errors if applied."""
diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py
index 805b8c86b0804..c93c59d1f4c42 100644
--- a/tests/v1/spec_decode/test_eagle.py
+++ b/tests/v1/spec_decode/test_eagle.py
@@ -3,7 +3,6 @@
from unittest import mock
-import numpy as np
import pytest
import torch
@@ -113,9 +112,7 @@ def test_prepare_next_token_ids():
sampled_token_ids_tensor = torch.tensor(
sampled_token_ids, dtype=torch.int32, device=device
)
- sampled_token_ids_cpu = [
- np.array([i for i in seq if i != -1]) for seq in sampled_token_ids
- ]
+ sampled_token_ids_cpu = [[i for i in seq if i != -1] for seq in sampled_token_ids]
expected_next_token_ids_cpu = [1, 4, 30, 40]
expected_next_token_ids_tensor = torch.tensor(
diff --git a/tests/v1/spec_decode/test_ngram.py b/tests/v1/spec_decode/test_ngram.py
index 563bc1d957f41..692c39282c372 100644
--- a/tests/v1/spec_decode/test_ngram.py
+++ b/tests/v1/spec_decode/test_ngram.py
@@ -77,7 +77,7 @@ def test_ngram_proposer():
# No match.
token_ids_cpu = np.array([[1, 2, 3, 4, 5]])
result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -88,7 +88,7 @@ def test_ngram_proposer():
# No match for 4-gram.
token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])
result = get_ngram_proposer(min_n=4, max_n=4, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -99,7 +99,7 @@ def test_ngram_proposer():
# No match for 4-gram but match for 3-gram.
token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])
result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -111,7 +111,7 @@ def test_ngram_proposer():
# In this case, the proposer should return the 4-gram match.
token_ids_cpu = np.array([[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]])
result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -122,7 +122,7 @@ def test_ngram_proposer():
# Match for 2-gram and 3-gram, but not 4-gram.
token_ids_cpu = np.array([[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]])
result = get_ngram_proposer(min_n=2, max_n=4, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -133,7 +133,7 @@ def test_ngram_proposer():
# Multiple 3-gram matched, but always pick the first one.
token_ids_cpu = np.array([[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]])
result = get_ngram_proposer(min_n=3, max_n=3, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -144,7 +144,7 @@ def test_ngram_proposer():
# check empty input
token_ids_cpu = np.array([[]])
result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
- sampled_token_ids=[np.array([0])],
+ sampled_token_ids=[[0]],
req_ids=["0"],
num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
token_ids_cpu=token_ids_cpu,
@@ -157,7 +157,7 @@ def test_ngram_proposer():
# second request has 3 tokens and no match. Padded with -1 for max len 5
token_ids_cpu = np.array([[1, 2, 3, 1, 2], [4, 5, 6, -1, -1]])
result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
- sampled_token_ids=[np.array([0]), np.array([1])],
+ sampled_token_ids=[[0], [1]],
req_ids=["0", "1"],
num_tokens_no_spec=np.array([5, 3]),
token_ids_cpu=token_ids_cpu,
@@ -181,7 +181,7 @@ def test_ngram_proposer():
input_2[:3] = [4, 5, 6]
token_ids_cpu = np.array([input_1, input_2])
result = ngram_proposer.propose(
- sampled_token_ids=[np.array([0]), np.array([1])],
+ sampled_token_ids=[[0], [1]],
req_ids=["0", "1"],
num_tokens_no_spec=np.array([len(input_1), 3]),
token_ids_cpu=token_ids_cpu,
diff --git a/tests/v1/spec_decode/test_tree_attention.py b/tests/v1/spec_decode/test_tree_attention.py
index 6958d62dc7e90..a4ee53008ce82 100644
--- a/tests/v1/spec_decode/test_tree_attention.py
+++ b/tests/v1/spec_decode/test_tree_attention.py
@@ -3,6 +3,7 @@
import math
+import pytest
import torch
from tests.v1.attention.utils import (
@@ -11,9 +12,16 @@ from tests.v1.attention.utils import (
try_get_attention_backend,
)
from vllm.attention.backends.registry import AttentionBackendEnum
+from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
from vllm.config import ParallelConfig, SpeculativeConfig
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+if not is_flash_attn_varlen_func_available():
+ pytest.skip(
+ "This test requires flash_attn_varlen_func, but it's not available.",
+ allow_module_level=True,
+ )
+
class MockAttentionLayer(torch.nn.Module):
_q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py
index b95c8df3469b3..d0f1b703fcb92 100644
--- a/tests/v1/worker/test_gpu_model_runner.py
+++ b/tests/v1/worker/test_gpu_model_runner.py
@@ -185,7 +185,9 @@ def _make_mock_backend_for_kernel_block_size(
supported_sizes: list[int | MultipleOf],
):
class _MockBackend:
- supported_kernel_block_sizes = supported_sizes
+ @staticmethod
+ def get_supported_kernel_block_sizes():
+ return supported_sizes
return _MockBackend()
@@ -483,7 +485,10 @@ def test_kv_cache_stride_order(monkeypatch, model_runner):
# Permutation that gets you back to expected kv shape
for test_stride in ((1, 4, 0, 2, 3), (0, 1, 2, 3, 4)):
- def rnd_stride_order(test_stride=test_stride):
+ def rnd_stride_order(
+ include_num_layers_dimension: bool = False, test_stride=test_stride
+ ):
+ assert not include_num_layers_dimension
return test_stride
# Patch the attention backend class and re-trigger the KV cache creation
@@ -956,7 +961,7 @@ def test_hybrid_block_table_initialization():
max_num_reqs = 10
max_num_blocks_per_req = 20
max_num_batched_tokens = 512
- dcp_kv_cache_interleave_size = 8
+ cp_kv_cache_interleave_size = 8
block_table = BlockTable(
block_size=block_size,
@@ -966,7 +971,7 @@ def test_hybrid_block_table_initialization():
pin_memory=False,
device=torch.device(DEVICE),
kernel_block_size=kernel_block_sizes[0],
- dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
)
# Verify hybrid block configuration
diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py
new file mode 100644
index 0000000000000..f7255fae05a4e
--- /dev/null
+++ b/tests/v1/worker/test_gpu_profiler.py
@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+
+import vllm.envs as envs
+from vllm.profiler.gpu_profiler import WorkerProfiler
+
+
+class ConcreteWorkerProfiler(WorkerProfiler):
+ """
+ A basic implementation of a worker profiler for testing purposes.
+ """
+
+ def __init__(self):
+ self.start_call_count = 0
+ self.stop_call_count = 0
+ self.should_fail_start = False
+ super().__init__()
+
+ def _start(self) -> None:
+ if self.should_fail_start:
+ raise RuntimeError("Simulated start failure")
+ self.start_call_count += 1
+
+ def _stop(self) -> None:
+ self.stop_call_count += 1
+
+
+@pytest.fixture(autouse=True)
+def reset_mocks():
+ """Fixture to reset mocks and env variables before each test."""
+ envs.VLLM_PROFILER_DELAY_ITERS = 0
+ envs.VLLM_PROFILER_MAX_ITERS = 0
+
+
+def test_immediate_start_stop():
+ """Test standard start without delay."""
+ profiler = ConcreteWorkerProfiler()
+
+ profiler.start()
+ assert profiler._running is True
+ assert profiler._active is True
+ assert profiler.start_call_count == 1
+
+ profiler.stop()
+ assert profiler._running is False
+ assert profiler._active is False
+ assert profiler.stop_call_count == 1
+
+
+def test_delayed_start():
+ """Test that profiler waits for N steps before actually starting."""
+ envs.VLLM_PROFILER_DELAY_ITERS = 2
+ profiler = ConcreteWorkerProfiler()
+
+ # User requests start
+ profiler.start()
+
+ # Should be active (request accepted) but not running (waiting for delay)
+ assert profiler._active is True
+ assert profiler._running is False
+ assert profiler.start_call_count == 0
+
+ # Step 1
+ profiler.step()
+ assert profiler._running is False
+
+ # Step 2 (Threshold reached)
+ profiler.step()
+ assert profiler._running is True
+ assert profiler.start_call_count == 1
+
+
+def test_max_iterations():
+ """Test that profiler stops automatically after max iterations."""
+ envs.VLLM_PROFILER_MAX_ITERS = 2
+ profiler = ConcreteWorkerProfiler()
+
+ profiler.start()
+ assert profiler._running is True
+
+ # Iteration 1
+ profiler.step() # profiling_count becomes 1
+ assert profiler._running is True
+
+ # Iteration 2
+ profiler.step() # profiling_count becomes 2
+ assert profiler._running is True
+
+ # Iteration 3 (Exceeds max)
+ profiler.step() # profiling_count becomes 3
+
+ # Should have stopped now
+ assert profiler._running is False
+ assert profiler.stop_call_count == 1
+
+
+def test_delayed_start_and_max_iters():
+ """Test combined delayed start and max iterations."""
+ envs.VLLM_PROFILER_DELAY_ITERS = 2
+ envs.VLLM_PROFILER_MAX_ITERS = 2
+ profiler = ConcreteWorkerProfiler()
+
+ profiler.start()
+
+ # Step 1
+ profiler.step()
+ assert profiler._running is False
+ assert profiler._active is True
+
+ # Step 2 (Starts now)
+ profiler.step()
+ assert profiler._profiling_for_iters == 1
+ assert profiler._running is True
+ assert profiler._active is True
+
+ # Next iteration
+ profiler.step()
+ assert profiler._profiling_for_iters == 2
+ assert profiler._running is True
+
+ # Iteration 2 (exceeds max)
+ profiler.step()
+
+ # Should have stopped now
+ assert profiler._running is False
+ assert profiler.stop_call_count == 1
+
+
+def test_idempotency():
+ """Test that calling start/stop multiple times doesn't break logic."""
+ profiler = ConcreteWorkerProfiler()
+
+ # Double Start
+ profiler.start()
+ profiler.start()
+ assert profiler.start_call_count == 1 # Should only start once
+
+ # Double Stop
+ profiler.stop()
+ profiler.stop()
+ assert profiler.stop_call_count == 1 # Should only stop once
+
+
+def test_step_inactive():
+ """Test that stepping while inactive does nothing."""
+ envs.VLLM_PROFILER_DELAY_ITERS = 2
+ profiler = ConcreteWorkerProfiler()
+
+ # Not started yet
+ profiler.step()
+ profiler.step()
+
+ # Even though we stepped 2 times, start shouldn't happen because active=False
+ assert profiler.start_call_count == 0
+
+
+def test_start_failure():
+ """Test behavior when the underlying _start method raises exception."""
+ profiler = ConcreteWorkerProfiler()
+ profiler.should_fail_start = True
+
+ profiler.start()
+
+ # Exception caught in _call_start
+ assert profiler._running is False # Should not mark as running
+ assert profiler._active is True # Request is still considered active
+ assert profiler.start_call_count == 0 # Logic failed inside start
+
+
+def test_shutdown():
+ """Test that shutdown calls stop only if running."""
+ profiler = ConcreteWorkerProfiler()
+
+ # Case 1: Not running
+ profiler.shutdown()
+ assert profiler.stop_call_count == 0
+
+ # Case 2: Running
+ profiler.start()
+ profiler.shutdown()
+ assert profiler.stop_call_count == 1
+
+
+def test_mixed_delay_and_stop():
+ """Test manual stop during the delay period."""
+ envs.VLLM_PROFILER_DELAY_ITERS = 5
+ profiler = ConcreteWorkerProfiler()
+
+ profiler.start()
+ profiler.step()
+ profiler.step()
+
+ # User cancels before delay finishes
+ profiler.stop()
+ assert profiler._active is False
+
+ # Further steps should not trigger start
+ profiler.step()
+ profiler.step()
+ profiler.step()
+
+ assert profiler.start_call_count == 0
diff --git a/tests/weight_loading/models-amd.txt b/tests/weight_loading/models-amd.txt
new file mode 100644
index 0000000000000..e31e904c08af4
--- /dev/null
+++ b/tests/weight_loading/models-amd.txt
@@ -0,0 +1,3 @@
+fp8, amd/Meta-Llama-3.1-8B-Instruct-FP8-KV, main
+None, amd/Llama-3.2-1B-Instruct-FP8-KV, main
+fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main
diff --git a/tests/weight_loading/models-large-amd.txt b/tests/weight_loading/models-large-amd.txt
new file mode 100644
index 0000000000000..b6f5b4b16b37f
--- /dev/null
+++ b/tests/weight_loading/models-large-amd.txt
@@ -0,0 +1,3 @@
+fp8, amd/Meta-Llama-3.1-70B-Instruct-FP8-KV, main
+None, microsoft/phi-4, main
+fp8, amd/Mixtral-8x22B-Instruct-v0.1-FP8-KV, main
diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh
index 5ea543f4cb1e8..1cea1bef8dbc9 100755
--- a/tools/ep_kernels/install_python_libraries.sh
+++ b/tools/ep_kernels/install_python_libraries.sh
@@ -1,94 +1,79 @@
#!/usr/bin/env bash
set -ex
-# prepare workspace directory
-WORKSPACE=$1
-if [ -z "$WORKSPACE" ]; then
- export WORKSPACE=$(pwd)/ep_kernels_workspace
-fi
+# usage: ./build.sh [workspace_dir] [mode]
+# mode: "install" (default) → install directly into current Python env
+# "wheel" → build wheels into WORKSPACE/dist
-if [ ! -d "$WORKSPACE" ]; then
- mkdir -p $WORKSPACE
-fi
+WORKSPACE=${1:-$(pwd)/ep_kernels_workspace}
+MODE=${2:-install}
+mkdir -p "$WORKSPACE"
+
+WHEEL_DIR="$WORKSPACE/dist"
+mkdir -p "$WHEEL_DIR"
+NVSHMEM_VER=3.3.9
+
+pushd "$WORKSPACE"
-# configurable pip command (default: pip3)
-PIP_CMD=${PIP_CMD:-pip3}
CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
# install dependencies if not installed
-$PIP_CMD install cmake torch ninja
-
-# build nvshmem
-pushd $WORKSPACE
-mkdir -p nvshmem_src
-wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.2.5/source/nvshmem_src_3.2.5-1.txz
-tar -xvf nvshmem_src_3.2.5-1.txz -C nvshmem_src --strip-components=1
-pushd nvshmem_src
-wget https://github.com/deepseek-ai/DeepEP/raw/main/third-party/nvshmem.patch
-git init
-git apply -vvv nvshmem.patch
-
-# assume CUDA_HOME is set correctly
-if [ -z "$CUDA_HOME" ]; then
- echo "CUDA_HOME is not set, please set it to your CUDA installation directory."
- exit 1
+if [ -z "$VIRTUAL_ENV" ]; then
+ uv pip install --system cmake torch ninja
+else
+ uv pip install cmake torch ninja
fi
-# assume TORCH_CUDA_ARCH_LIST is set correctly
-if [ -z "$TORCH_CUDA_ARCH_LIST" ]; then
- echo "TORCH_CUDA_ARCH_LIST is not set, please set it to your desired architecture."
+# fetch nvshmem
+ARCH=$(uname -m)
+case "${ARCH,,}" in
+ x86_64|amd64)
+ NVSHMEM_SUBDIR="linux-x86_64"
+ NVSHMEM_FILE="libnvshmem-linux-x86_64-${NVSHMEM_VER}_cuda12-archive.tar.xz"
+ ;;
+ aarch64|arm64)
+ NVSHMEM_SUBDIR="linux-sbsa"
+ NVSHMEM_FILE="libnvshmem-linux-sbsa-${NVSHMEM_VER}_cuda12-archive.tar.xz"
+ ;;
+ *)
+ echo "Unsupported architecture: ${ARCH}" >&2
exit 1
-fi
+ ;;
+esac
-# disable all features except IBGDA
-export NVSHMEM_IBGDA_SUPPORT=1
-
-export NVSHMEM_SHMEM_SUPPORT=0
-export NVSHMEM_UCX_SUPPORT=0
-export NVSHMEM_USE_NCCL=0
-export NVSHMEM_PMIX_SUPPORT=0
-export NVSHMEM_TIMEOUT_DEVICE_POLLING=0
-export NVSHMEM_USE_GDRCOPY=0
-export NVSHMEM_IBRC_SUPPORT=0
-export NVSHMEM_BUILD_TESTS=0
-export NVSHMEM_BUILD_EXAMPLES=0
-export NVSHMEM_MPI_SUPPORT=0
-export NVSHMEM_BUILD_HYDRA_LAUNCHER=0
-export NVSHMEM_BUILD_TXZ_PACKAGE=0
-export NVSHMEM_TIMEOUT_DEVICE_POLLING=0
-
-cmake -G Ninja -S . -B $WORKSPACE/nvshmem_build/ -DCMAKE_INSTALL_PREFIX=$WORKSPACE/nvshmem_install
-cmake --build $WORKSPACE/nvshmem_build/ --target install
+NVSHMEM_URL="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${NVSHMEM_SUBDIR}/${NVSHMEM_FILE}"
+pushd "$WORKSPACE"
+echo "Downloading NVSHMEM ${NVSHMEM_VER} for ${NVSHMEM_SUBDIR} ..."
+curl -fSL "${NVSHMEM_URL}" -o "${NVSHMEM_FILE}"
+tar -xf "${NVSHMEM_FILE}"
+mv "${NVSHMEM_FILE%.tar.xz}" nvshmem
+rm -f "${NVSHMEM_FILE}"
+rm -rf nvshmem/lib/bin nvshmem/lib/share
popd
-export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem_install:$CMAKE_PREFIX_PATH
+export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem/lib/cmake:$CMAKE_PREFIX_PATH
is_git_dirty() {
local dir=$1
pushd "$dir" > /dev/null
-
- if [ -d ".git" ] && [ -n "$(git status --porcelain 2>/dev/null)" ]; then
+ if [ -d ".git" ] && [ -n "$(git status --porcelain 3>/dev/null)" ]; then
popd > /dev/null
- return 0 # dirty (true)
+ return 0
else
popd > /dev/null
- return 1 # clean (false)
+ return 1
fi
}
-# Function to handle git repository cloning with dirty/incomplete checks
clone_repo() {
local repo_url=$1
local dir_name=$2
local key_file=$3
local commit_hash=$4
-
if [ -d "$dir_name" ]; then
- # Check if directory has uncommitted changes (dirty)
if is_git_dirty "$dir_name"; then
echo "$dir_name directory is dirty, skipping clone"
- # Check if clone failed (directory exists but not a valid git repo or missing key files)
elif [ ! -d "$dir_name/.git" ] || [ ! -f "$dir_name/$key_file" ]; then
echo "$dir_name directory exists but clone appears incomplete, cleaning up and re-cloning"
rm -rf "$dir_name"
@@ -99,7 +84,7 @@ clone_repo() {
cd ..
fi
else
- echo "$dir_name directory exists and appears complete; manually update if needed"
+ echo "$dir_name directory exists and appears complete"
fi
else
git clone "$repo_url"
@@ -111,17 +96,55 @@ clone_repo() {
fi
}
-# build and install pplx, require pytorch installed
-pushd $WORKSPACE
-clone_repo "https://github.com/ppl-ai/pplx-kernels" "pplx-kernels" "setup.py" "c336faf"
-cd pplx-kernels
-$PIP_CMD install --no-build-isolation -vvv -e .
-popd
+deepep_cuda13_patch() {
+ cuda_version_major=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
+ if [ ${cuda_version_major} -ge 13 ]; then
+ sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py"
+ fi
+}
-# build and install deepep, require pytorch installed
-pushd $WORKSPACE
-clone_repo "https://github.com/deepseek-ai/DeepEP" "DeepEP" "setup.py" "73b6ea4"
-cd DeepEP
-export NVSHMEM_DIR=$WORKSPACE/nvshmem_install
-$PIP_CMD install --no-build-isolation -vvv -e .
-popd
+do_build() {
+ local repo=$1
+ local name=$2
+ local key=$3
+ local commit=$4
+ local extra_env=$5
+
+ pushd "$WORKSPACE"
+ clone_repo "$repo" "$name" "$key" "$commit"
+ cd "$name"
+
+ if [ "$name" == "DeepEP" ]; then
+ deepep_cuda13_patch
+ fi
+
+ if [ "$MODE" = "install" ]; then
+ echo "Installing $name into environment"
+ eval "$extra_env" uv pip install --no-build-isolation -vvv .
+ else
+ echo "Building $name wheel into $WHEEL_DIR"
+ eval "$extra_env" uv build --wheel --no-build-isolation -vvv --out-dir "$WHEEL_DIR" .
+ fi
+ popd
+}
+
+# build pplx-kernels
+do_build \
+ "https://github.com/ppl-ai/pplx-kernels" \
+ "pplx-kernels" \
+ "setup.py" \
+ "12cecfd" \
+ ""
+
+# build DeepEP
+do_build \
+ "https://github.com/deepseek-ai/DeepEP" \
+ "DeepEP" \
+ "setup.py" \
+ "73b6ea4" \
+ "export NVSHMEM_DIR=$WORKSPACE/nvshmem; "
+
+if [ "$MODE" = "wheel" ]; then
+ echo "All wheels written to $WHEEL_DIR"
+ ls -l "$WHEEL_DIR"
+fi
diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh
index 4f2cd302c3eff..ee9a5dd4aa643 100755
--- a/tools/install_deepgemm.sh
+++ b/tools/install_deepgemm.sh
@@ -1,12 +1,13 @@
#!/bin/bash
-# Script to install DeepGEMM from source
-# This script can be used both in Docker builds and by users locally
-
+# Script to build and/or install DeepGEMM from source
+# Default: build and install immediately
+# Optional: build wheels to a directory for later installation (useful in multi-stage builds)
set -e
# Default values
DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
DEEPGEMM_GIT_REF="594953acce41793ae00a1233eb516044d604bcb6"
+WHEEL_DIR=""
# Parse command line arguments
while [[ $# -gt 0 ]]; do
@@ -27,11 +28,20 @@ while [[ $# -gt 0 ]]; do
CUDA_VERSION="$2"
shift 2
;;
+ --wheel-dir)
+ if [[ -z "$2" || "$2" =~ ^- ]]; then
+ echo "Error: --wheel-dir requires a directory path." >&2
+ exit 1
+ fi
+ WHEEL_DIR="$2"
+ shift 2
+ ;;
-h|--help)
echo "Usage: $0 [OPTIONS]"
echo "Options:"
echo " --ref REF Git reference to checkout (default: $DEEPGEMM_GIT_REF)"
echo " --cuda-version VER CUDA version (auto-detected if not provided)"
+ echo " --wheel-dir PATH If set, build wheel into PATH but do not install"
echo " -h, --help Show this help message"
exit 0
;;
@@ -57,16 +67,15 @@ fi
CUDA_MAJOR="${CUDA_VERSION%%.*}"
CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
CUDA_MINOR="${CUDA_MINOR%%.*}"
-
echo "CUDA version: $CUDA_VERSION (major: $CUDA_MAJOR, minor: $CUDA_MINOR)"
# Check CUDA version requirement
if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then
- echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
+ echo "Skipping DeepGEMM build/installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
exit 0
fi
-echo "Installing DeepGEMM from source..."
+echo "Preparing DeepGEMM build..."
echo "Repository: $DEEPGEMM_GIT_REPO"
echo "Reference: $DEEPGEMM_GIT_REF"
@@ -76,23 +85,31 @@ trap 'rm -rf "$INSTALL_DIR"' EXIT
# Clone the repository
git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm"
-
-echo "🏗️ Building DeepGEMM"
pushd "$INSTALL_DIR/deepgemm"
# Checkout the specific reference
git checkout "$DEEPGEMM_GIT_REF"
-# Build DeepGEMM
+# Clean previous build artifacts
# (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
-rm -rf build dist
-rm -rf *.egg-info
+rm -rf build dist *.egg-info
+
+# Build wheel
+echo "🏗️ Building DeepGEMM wheel..."
python3 setup.py bdist_wheel
-# Install the wheel
+# If --wheel-dir was specified, copy wheels there and exit
+if [ -n "$WHEEL_DIR" ]; then
+ mkdir -p "$WHEEL_DIR"
+ cp dist/*.whl "$WHEEL_DIR"/
+ echo "✅ Wheel built and copied to $WHEEL_DIR"
+ popd
+ exit 0
+fi
+
+# Default behaviour: install built wheel
if command -v uv >/dev/null 2>&1; then
echo "Installing DeepGEMM wheel using uv..."
- # Use --system in Docker contexts, respect user's environment otherwise
if [ -n "$VLLM_DOCKER_BUILD_CONTEXT" ]; then
uv pip install --system dist/*.whl
else
@@ -104,5 +121,4 @@ else
fi
popd
-
echo "✅ DeepGEMM installation completed successfully"
diff --git a/tools/install_nixl_from_source_ubuntu.py b/tools/install_nixl_from_source_ubuntu.py
index a786abba95ad9..b8a55c615426e 100644
--- a/tools/install_nixl_from_source_ubuntu.py
+++ b/tools/install_nixl_from_source_ubuntu.py
@@ -95,6 +95,7 @@ def install_system_dependencies():
"meson",
"libtool",
"libtool-bin",
+ "pkg-config",
]
run_command(["apt-get", "update"])
run_command(["apt-get", "install", "-y"] + apt_packages)
diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py
index 8d04848f8f780..34f6e8c928ffb 100755
--- a/tools/pre_commit/mypy.py
+++ b/tools/pre_commit/mypy.py
@@ -38,6 +38,7 @@ FILES = [
"vllm/usage",
"vllm/v1/core",
"vllm/v1/engine",
+ "vllm/v1/worker",
]
# After fixing errors resulting from changing follow_imports
@@ -62,7 +63,6 @@ SEPARATE_GROUPS = [
"vllm/v1/sample",
"vllm/v1/spec_decode",
"vllm/v1/structured_output",
- "vllm/v1/worker",
]
# TODO(woosuk): Include the code from Megatron and HuggingFace.
diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py
index e53e4ae6e5296..db79b3f5e8bcb 100644
--- a/vllm/_aiter_ops.py
+++ b/vllm/_aiter_ops.py
@@ -948,6 +948,31 @@ class rocm_aiter_ops:
(8192, 32768),
]
+ @staticmethod
+ def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool:
+ return (n, k) in [
+ (8192, 4096),
+ (1280, 8192),
+ (16384, 53248),
+ (106496, 16384),
+ (57344, 8192),
+ (8192, 2048),
+ (2560, 8192),
+ (10240, 8192),
+ (16384, 16384),
+ (8192, 28672),
+ (28672, 8192),
+ (18432, 16384),
+ (8192, 1024),
+ (7168, 8192),
+ (5120, 8192),
+ (8192, 8192),
+ (8192, 7168),
+ (14336, 8192),
+ (8192, 14336),
+ (8192, 3584),
+ ]
+
@staticmethod
def shuffle_weight(
self, tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
index 096266c9764e8..0f625a7945241 100644
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -328,10 +328,7 @@ def rotary_embedding(
def rms_norm(
out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float
) -> None:
- # TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input
- # If removed, also need to remove contiguous in MatcherRMSNorm
- input_contiguous = input.contiguous()
- torch.ops._C.rms_norm(out, input_contiguous, weight, epsilon)
+ torch.ops._C.rms_norm(out, input, weight, epsilon)
def fused_add_rms_norm(
@@ -2702,6 +2699,31 @@ def cpu_attention_with_kv_cache(
)
+def cpu_gemm_wna16(
+ input: torch.Tensor,
+ q_weight: torch.Tensor,
+ scales: torch.Tensor,
+ zeros: torch.Tensor | None,
+ g_idx: torch.Tensor | None,
+ bias: torch.Tensor | None,
+ pack_factor: int,
+ isa_hint: str,
+) -> torch.Tensor:
+ output = torch.empty((input.size(0), scales.size(1)), dtype=input.dtype)
+ torch.ops._C.cpu_gemm_wna16(
+ input,
+ q_weight,
+ output,
+ scales,
+ zeros,
+ g_idx,
+ bias,
+ pack_factor,
+ isa_hint,
+ )
+ return output
+
+
if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
@register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
diff --git a/vllm/attention/__init__.py b/vllm/attention/__init__.py
index dd35165d5415e..8b4dc4013362e 100644
--- a/vllm/attention/__init__.py
+++ b/vllm/attention/__init__.py
@@ -7,7 +7,7 @@ from vllm.attention.backends.abstract import (
AttentionType,
)
from vllm.attention.layer import Attention
-from vllm.attention.selector import get_attn_backend
+from vllm.attention.selector import get_attn_backend, get_mamba_attn_backend
__all__ = [
"Attention",
@@ -15,4 +15,5 @@ __all__ = [
"AttentionMetadata",
"AttentionType",
"get_attn_backend",
+ "get_mamba_attn_backend",
]
diff --git a/vllm/attention/backends/abstract.py b/vllm/attention/backends/abstract.py
index 9275d70fd86a4..bd7e81b15bfc3 100644
--- a/vllm/attention/backends/abstract.py
+++ b/vllm/attention/backends/abstract.py
@@ -46,9 +46,12 @@ class AttentionBackend(ABC):
# makes sure the output tensor is allocated inside the cudagraph.
accept_output_buffer: bool = False
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(1)]
supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = ["auto"]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [MultipleOf(1)]
+
@staticmethod
@abstractmethod
def get_name() -> str:
@@ -76,7 +79,34 @@ class AttentionBackend(ABC):
raise NotImplementedError
@staticmethod
- def get_kv_cache_stride_order() -> tuple[int, ...]:
+ def get_kv_cache_stride_order(
+ include_num_layers_dimension: bool = False,
+ ) -> tuple[int, ...]:
+ """
+ Get the physical (memory layout) ordering of the kv cache dimensions.
+ e.g. if the KV cache shape is
+ [2, num_blocks, block_size, num_heads, head_size],
+ and get_kv_cache_stride_order returns (1, 3, 0, 2, 4) then the physical
+ ordering of dimensions is
+ [num_blocks, num_heads, 2, block_size, head_size].
+
+ If this function is unimplemented / raises NotImplementedError,
+ the physical layout of the KV cache will match the logical shape.
+
+ Args:
+ include_num_layers_dimension: if True, includes an additional
+ num_layers dimension, which is assumed to be prepended
+ to the logical KV cache shape.
+ With the above example, a return value (2, 4, 0, 1, 3, 5)
+ corresponds to
+ [num_blocks, num_heads, num_layers, 2, block_size, head_size].
+
+ If an additional dimension is NOT included in the returned
+ tuple, the physical layout will not include a layers dimension.
+
+ Returns:
+ A tuple of ints which is a permutation of range(len(shape)).
+ """
raise NotImplementedError
@classmethod
@@ -115,18 +145,17 @@ class AttentionBackend(ABC):
if block_size not in valid_sizes:
return False
- if not cls.supported_kernel_block_sizes:
+ supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()
+ if not supported_kernel_block_sizes:
return True
- for supported_size in cls.supported_kernel_block_sizes:
- is_multiple_of = (
- isinstance(supported_size, MultipleOf)
- and block_size % supported_size.base == 0
- )
- is_int_equal = (
- isinstance(supported_size, int) and block_size == supported_size
- )
- if is_multiple_of or is_int_equal:
+ for supported_size in supported_kernel_block_sizes:
+ if isinstance(supported_size, MultipleOf):
+ supported_size = supported_size.base
+ # With hybrid_blocks feature, the framework-level block size
+ # only needs to be a multiple of the kernel's requirement,
+ # even if the kernel requires a fixed block_size.
+ if block_size % supported_size == 0:
return True
return False
@@ -266,6 +295,12 @@ class AttentionImpl(ABC, Generic[T]):
dcp_world_size: int
dcp_rank: int
+ pcp_world_size: int
+ pcp_rank: int
+
+ total_cp_world_size: int
+ total_cp_rank: int
+
def __new__(cls, *args, **kwargs):
# use __new__ so that all subclasses will call this
self = super().__new__(cls)
@@ -278,6 +313,17 @@ class AttentionImpl(ABC, Generic[T]):
# DCP might not be initialized in testing
self.dcp_world_size = 1
self.dcp_rank = 0
+ try:
+ from vllm.distributed.parallel_state import get_pcp_group
+
+ self.pcp_world_size = get_pcp_group().world_size
+ self.pcp_rank = get_pcp_group().rank_in_group
+ except AssertionError:
+ self.pcp_world_size = 1
+ self.pcp_rank = 0
+ self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size
+ self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
+
self.need_to_return_lse_for_decode = (
self.dcp_world_size > 1 and self.can_return_lse_for_decode
)
diff --git a/vllm/attention/backends/registry.py b/vllm/attention/backends/registry.py
index f07a6059be377..125e4e3827747 100644
--- a/vllm/attention/backends/registry.py
+++ b/vllm/attention/backends/registry.py
@@ -2,8 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Attention backend registry"""
-import enum
from collections.abc import Callable
+from enum import Enum, EnumMeta
from typing import TYPE_CHECKING, cast
from vllm.logger import init_logger
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
-class _AttentionBackendEnumMeta(enum.EnumMeta):
+class _AttentionBackendEnumMeta(EnumMeta):
"""Metaclass for AttentionBackendEnum to provide better error messages."""
def __getitem__(cls, name: str):
@@ -23,15 +23,15 @@ class _AttentionBackendEnumMeta(enum.EnumMeta):
try:
return super().__getitem__(name)
except KeyError:
- members = cast("dict[str, AttentionBackendEnum]", cls.__members__).values()
- valid_backends = ", ".join(m.name for m in members)
+ members = cast("dict[str, Enum]", cls.__members__).keys()
+ valid_backends = ", ".join(members)
raise ValueError(
f"Unknown attention backend: '{name}'. "
f"Valid options are: {valid_backends}"
) from None
-class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
+class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
"""Enumeration of all supported attention backends.
The enum value is the default class path, but this can be overridden
@@ -43,12 +43,17 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
- XFORMERS = "vllm.v1.attention.backends.xformers.XFormersAttentionBackend"
ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
+ ROCM_AITER_TRITON_MLA = (
+ "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend"
+ )
ROCM_AITER_FA = (
"vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
)
+ ROCM_AITER_MLA_SPARSE = (
+ "vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse.ROCMAiterMLASparseBackend"
+ )
TORCH_SDPA = "" # this tag is only used for ViT
FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
FLASHINFER_MLA = (
@@ -83,7 +88,7 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
Raises:
ValueError: If Backend.CUSTOM is used without being registered
"""
- path = _OVERRIDES.get(self, self.value)
+ path = _ATTN_OVERRIDES.get(self, self.value)
if not path:
raise ValueError(
f"Backend {self.name} must be registered before use. "
@@ -111,18 +116,93 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
Returns:
True if the backend has a registered override
"""
- return self in _OVERRIDES
+ return self in _ATTN_OVERRIDES
def clear_override(self) -> None:
"""Clear any override for this backend, reverting to the default."""
- _OVERRIDES.pop(self, None)
+ _ATTN_OVERRIDES.pop(self, None)
-_OVERRIDES: dict[AttentionBackendEnum, str] = {}
+class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
+ """Enumeration of all supported mamba attention backends.
+
+ The enum value is the default class path, but this can be overridden
+ at runtime using register_backend().
+
+ To get the actual backend class (respecting overrides), use:
+ backend.get_class()
+ """
+
+ MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend"
+ MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend"
+ SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend"
+ LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend"
+ GDN_ATTN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend"
+ # Placeholder for third-party/custom backends - must be registered before use
+ CUSTOM = ""
+
+ def get_path(self, include_classname: bool = True) -> str:
+ """Get the class path for this backend (respects overrides).
+
+ Returns:
+ The fully qualified class path string
+
+ Raises:
+ ValueError: If Backend.CUSTOM is used without being registered
+ """
+ path = _MAMBA_ATTN_OVERRIDES.get(self, self.value)
+ if not path:
+ raise ValueError(
+ f"Backend {self.name} must be registered before use. "
+ f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')"
+ )
+ if not include_classname:
+ path = path.rsplit(".", 1)[0]
+ return path
+
+ def get_class(self) -> "type[AttentionBackend]":
+ """Get the backend class (respects overrides).
+
+ Returns:
+ The backend class
+
+ Raises:
+ ImportError: If the backend class cannot be imported
+ ValueError: If Backend.CUSTOM is used without being registered
+ """
+ return resolve_obj_by_qualname(self.get_path())
+
+ def is_overridden(self) -> bool:
+ """Check if this backend has been overridden.
+
+ Returns:
+ True if the backend has a registered override
+ """
+ return self in _MAMBA_ATTN_OVERRIDES
+
+ def clear_override(self) -> None:
+ """Clear any override for this backend, reverting to the default."""
+ _MAMBA_ATTN_OVERRIDES.pop(self, None)
+
+
+MAMBA_TYPE_TO_BACKEND_MAP = {
+ "mamba1": MambaAttentionBackendEnum.MAMBA1.name,
+ "mamba2": MambaAttentionBackendEnum.MAMBA2.name,
+ "short_conv": MambaAttentionBackendEnum.SHORT_CONV.name,
+ "linear_attention": MambaAttentionBackendEnum.LINEAR.name,
+ "gdn_attention": MambaAttentionBackendEnum.GDN_ATTN.name,
+ "custom": MambaAttentionBackendEnum.CUSTOM.name,
+}
+
+
+_ATTN_OVERRIDES: dict[AttentionBackendEnum, str] = {}
+_MAMBA_ATTN_OVERRIDES: dict[MambaAttentionBackendEnum, str] = {}
def register_backend(
- backend: AttentionBackendEnum, class_path: str | None = None
+ backend: AttentionBackendEnum | MambaAttentionBackendEnum,
+ is_mamba: bool = False,
+ class_path: str | None = None,
) -> Callable[[type], type]:
"""Register or override a backend implementation.
@@ -135,12 +215,17 @@ def register_backend(
Decorator function if class_path is None, otherwise a no-op
Examples:
- # Override an existing backend
+ # Override an existing attention backend
@register_backend(AttentionBackendEnum.FLASH_ATTN)
class MyCustomFlashAttn:
...
- # Register a custom third-party backend
+ # Override an existing mamba attention backend
+ @register_backend(MambaAttentionBackendEnum.LINEAR, is_mamba=True)
+ class MyCustomMambaAttn:
+ ...
+
+ # Register a custom third-party attention backend
@register_backend(AttentionBackendEnum.CUSTOM)
class MyCustomBackend:
...
@@ -153,11 +238,17 @@ def register_backend(
"""
def decorator(cls: type) -> type:
- _OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
+ if is_mamba:
+ _MAMBA_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index]
+ else:
+ _ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index]
return cls
if class_path is not None:
- _OVERRIDES[backend] = class_path
+ if is_mamba:
+ _MAMBA_ATTN_OVERRIDES[backend] = class_path # type: ignore[index]
+ else:
+ _ATTN_OVERRIDES[backend] = class_path # type: ignore[index]
return lambda x: x
return decorator
diff --git a/vllm/attention/layer.py b/vllm/attention/layer.py
index 37f9a4b383ce9..f1d57ac50fb9f 100644
--- a/vllm/attention/layer.py
+++ b/vllm/attention/layer.py
@@ -51,31 +51,6 @@ else:
FP8_DTYPE = current_platform.fp8_dtype()
logger = init_logger(__name__)
-USE_XFORMERS_OPS = None
-
-
-def check_xformers_availability():
- global USE_XFORMERS_OPS
- if USE_XFORMERS_OPS is not None:
- return USE_XFORMERS_OPS
-
- if current_platform.is_cuda() and current_platform.has_device_capability(100):
- # Xformers FA is not compatible with B200
- USE_XFORMERS_OPS = False
- else:
- try:
- from importlib.util import find_spec
-
- find_spec("xformers.ops")
- USE_XFORMERS_OPS = True
- except ImportError:
- USE_XFORMERS_OPS = False
-
- # the warning only needs to be shown once
- if not USE_XFORMERS_OPS:
- logger.warning("Xformers is not available, falling back.")
-
- return USE_XFORMERS_OPS
def check_upstream_fa_availability(dtype: torch.dtype):
@@ -310,7 +285,8 @@ class Attention(nn.Module, AttentionLayerBase):
kv_sharing_target_layer_name,
**extra_impl_args,
)
- self.backend = AttentionBackendEnum[self.attn_backend.get_name()]
+ backend_name = self.attn_backend.get_name()
+ self.backend = AttentionBackendEnum.__members__.get(backend_name)
self.dtype = dtype
# For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
@@ -532,7 +508,6 @@ class MultiHeadAttention(nn.Module):
if backend
in {
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.PALLAS,
AttentionBackendEnum.ROCM_AITER_FA,
AttentionBackendEnum.FLASH_ATTN,
@@ -548,12 +523,6 @@ class MultiHeadAttention(nn.Module):
)
)
- if (
- self.attn_backend == AttentionBackendEnum.XFORMERS
- and not check_xformers_availability()
- ):
- self.attn_backend = AttentionBackendEnum.TORCH_SDPA
-
self.is_flash_attn_backend = self.attn_backend in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.ROCM_AITER_FA,
@@ -613,12 +582,6 @@ class MultiHeadAttention(nn.Module):
max_seqlen_k=kv_len,
softmax_scale=self.scale,
)
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- from xformers import ops as xops
-
- out = xops.memory_efficient_attention_forward(
- query, key, value, scale=self.scale
- )
elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
query, key, value = (x.transpose(1, 2) for x in (query, key, value))
out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)
diff --git a/vllm/attention/ops/common.py b/vllm/attention/ops/common.py
index 2cbb5c91cc3b3..67c5f7dbba9c0 100644
--- a/vllm/attention/ops/common.py
+++ b/vllm/attention/ops/common.py
@@ -169,12 +169,11 @@ def correct_attn_out(
return out, lse
-def cp_lse_ag_out_rs(
+def _cp_lse_common(
cp_attn_out: torch.Tensor,
cp_attn_lse: torch.Tensor,
cp_group: GroupCoordinator,
- ctx: CPTritonContext = None,
- return_lse=False,
+ ctx: CPTritonContext | None = None,
):
"""
cp_attn_out: [ B, H, D ]
@@ -195,6 +194,22 @@ def cp_lse_ag_out_rs(
cp_attn_lse = cp_attn_lse.contiguous()
lses = cp_group.all_gather(cp_attn_lse, dim=0).view_as(lses)
out, lse = correct_attn_out(cp_attn_out, lses, cp_group.rank_in_group, ctx)
+ assert out.is_contiguous()
+ return out, lse
+
+
+def cp_lse_ag_out_rs(
+ cp_attn_out: torch.Tensor,
+ cp_attn_lse: torch.Tensor,
+ cp_group: GroupCoordinator,
+ ctx: CPTritonContext | None = None,
+ return_lse: bool = False,
+):
+ """
+ cp_attn_out: [ B, H, D ]
+ cp_attn_lse: [ B, H ]
+ """
+ out, lse = _cp_lse_common(cp_attn_out, cp_attn_lse, cp_group, ctx=ctx)
out = cp_group.reduce_scatter(out, dim=1)
if return_lse:
@@ -205,6 +220,25 @@ def cp_lse_ag_out_rs(
return out
+def cp_lse_ag_out_ar(
+ cp_attn_out: torch.Tensor,
+ cp_attn_lse: torch.Tensor,
+ cp_group: GroupCoordinator,
+ ctx: CPTritonContext | None = None,
+ return_lse: bool = False,
+):
+ """
+ cp_attn_out: [ B, H, D ]
+ cp_attn_lse: [ B, H ]
+ """
+ out, lse = _cp_lse_common(cp_attn_out, cp_attn_lse, cp_group, ctx=ctx)
+ out = cp_group.all_reduce(out)
+
+ if return_lse:
+ return out, lse
+ return out
+
+
@triton.jit
def _pack_seq_kernel(
x_ptr, # [N, D]
diff --git a/vllm/attention/ops/rocm_aiter_mla_sparse.py b/vllm/attention/ops/rocm_aiter_mla_sparse.py
new file mode 100644
index 0000000000000..080e92ecc9408
--- /dev/null
+++ b/vllm/attention/ops/rocm_aiter_mla_sparse.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import importlib
+from functools import lru_cache
+
+import torch
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+# Take from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L84
+def fp8_mqa_logits_torch(
+ q: torch.Tensor,
+ kv: tuple[torch.Tensor, torch.Tensor],
+ weights: torch.Tensor,
+ cu_seqlen_ks: torch.Tensor,
+ cu_seqlen_ke: torch.Tensor,
+) -> torch.Tensor:
+ """Compute FP8 MQA logits for a single sequence without KV paging.
+
+ Args:
+ q: Query tensor of shape [M, H, D]. Casted to
+ `torch.float8_e4m3fn` by caller.
+ kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
+ dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
+ [N, 1]) with dtype `torch.float32`.
+ weights: weights of shape [M, H], dtype `torch.float32`.
+ cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
+ shape [M], dtype int32.
+ cu_seqlen_ke: End indices (exclusive) for valid K per query position,
+ shape [M], dtype int32.
+
+ Returns:
+ Logits tensor of shape [M, N], dtype `torch.float32`.
+ """
+ kv, scale = kv
+ seq_len_kv = kv.shape[0]
+ k = kv.to(torch.bfloat16)
+ q = q.to(torch.bfloat16)
+
+ mask_lo = (
+ torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
+ )
+ mask_hi = (
+ torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
+ )
+ mask = mask_lo & mask_hi
+
+ score = torch.einsum("mhd,nd->hmn", q, k).float() * scale
+ logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
+ logits = logits.masked_fill(~mask, float("-inf"))
+
+ return logits
+
+
+def rocm_fp8_mqa_logits(
+ q: torch.Tensor,
+ kv: tuple[torch.Tensor, torch.Tensor],
+ weights: torch.Tensor,
+ cu_seqlen_ks: torch.Tensor,
+ cu_seqlen_ke: torch.Tensor,
+) -> torch.Tensor:
+ """Compute FP8 MQA logits for a single sequence without KV paging.
+
+ Args:
+ q: Query tensor of shape [M, H, D]. Casted to
+ `torch.float8_e4m3fn` by caller.
+ kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
+ dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
+ [N, 1]) with dtype `torch.float32`.
+ weights: weights of shape [M, H], dtype `torch.float32`.
+ cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
+ shape [M], dtype int32.
+ cu_seqlen_ke: End indices (exclusive) for valid K per query position,
+ shape [M], dtype int32.
+
+ Returns:
+ Logits tensor of shape [M, N], dtype `torch.float32`.
+ """
+
+ # TODO(ganyi): Temporarily workaround, will remove the module check and reference
+ # path after aiter merge this kernel into main
+ @lru_cache
+ def has_mqa_logits_module():
+ return importlib.util.find_spec("aiter.ops.triton.fp8_mqa_logits") is not None
+
+ if rocm_aiter_ops.is_enabled() and has_mqa_logits_module():
+ from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
+
+ kv, scale = kv
+ return fp8_mqa_logits(q, kv, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
+ else:
+ return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
+
+
+# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156
+def fp8_paged_mqa_logits_torch(
+ q: torch.Tensor,
+ kv_cache: torch.Tensor,
+ weights: torch.Tensor,
+ context_lens: torch.Tensor,
+ block_tables: torch.Tensor,
+ max_model_len: int,
+):
+ from vllm.utils.math_utils import cdiv
+
+ fp8_dtype = current_platform.fp8_dtype()
+ batch_size, next_n, _, dim = q.size()
+ kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
+ scale = scale.contiguous().view(torch.float)
+ q = q.float()
+ kv_cache = kv_cache.view(fp8_dtype).float() * scale
+ num_block, block_size, _, dim = kv_cache.size()
+ logits = torch.full(
+ [batch_size * next_n, max_model_len],
+ float("-inf"),
+ device=q.device,
+ dtype=torch.float32,
+ )
+ context_lens = context_lens.tolist()
+ for i in range(batch_size):
+ context_len = context_lens[i]
+ q_offsets = torch.arange(context_len - next_n, context_len, device="cuda")
+ weight_slice = (
+ weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
+ )
+ for block_rk in range(cdiv(context_len, block_size)):
+ block_idx = block_tables[i][block_rk]
+ qx, kx = q[i], kv_cache[block_idx]
+ k_offsets = torch.arange(
+ block_rk * block_size, (block_rk + 1) * block_size, device="cuda"
+ )
+ mask = (k_offsets[None, :] < context_len) & (
+ k_offsets[None, :] <= q_offsets[:, None]
+ )
+ s = torch.where(
+ mask[None, :, :],
+ (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
+ logits.dtype
+ ),
+ float("-inf"),
+ )
+ s = torch.relu(s) * weight_slice[..., None]
+ s = s.sum(dim=0)
+ logits[
+ i * next_n : (i + 1) * next_n,
+ block_rk * block_size : (block_rk + 1) * block_size,
+ ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf"))
+ return logits
+
+
+def rocm_fp8_paged_mqa_logits(
+ q_fp8: torch.Tensor,
+ kv_cache_fp8: torch.Tensor,
+ weights: torch.Tensor,
+ context_lens: torch.Tensor,
+ block_tables: torch.Tensor,
+ schedule_metadata: torch.Tensor,
+ max_model_len: int,
+) -> torch.Tensor:
+ """Compute FP8 MQA logits using paged KV-cache.
+
+ Args:
+ q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
+ `torch.float8_e4m3fn` by caller.
+ kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
+ [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
+ 4 bytes per (block,pos) store the `float` dequant scale.
+ weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
+ context_lens: Tensor of shape [B], dtype int32; effective context length
+ for each batch element.
+ block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
+ block indices to physical blocks in the paged cache.
+ schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
+ used to distribute work across SMs.
+ max_model_len: Maximum sequence length used to size the logits output.
+
+ Returns:
+ Logits tensor of shape [B * next_n, max_model_len], dtype
+ `torch.float32`.
+ """
+
+ if rocm_aiter_ops.is_enabled():
+ from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits_stage1
+
+ batch_size, next_n, heads, _ = q_fp8.shape
+ out_qk = torch.full(
+ (heads, batch_size * next_n, max_model_len),
+ float("-inf"),
+ device="cuda",
+ dtype=torch.float32,
+ )
+ deepgemm_fp8_paged_mqa_logits_stage1(
+ q_fp8,
+ kv_cache_fp8,
+ weights,
+ out_qk,
+ context_lens,
+ block_tables,
+ max_model_len,
+ )
+ return out_qk.sum(dim=0)
+ else:
+ return fp8_paged_mqa_logits_torch(
+ q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
+ )
diff --git a/vllm/attention/ops/vit_attn_wrappers.py b/vllm/attention/ops/vit_attn_wrappers.py
index 06a9f7cd82266..46f8f5117f7a7 100644
--- a/vllm/attention/ops/vit_attn_wrappers.py
+++ b/vllm/attention/ops/vit_attn_wrappers.py
@@ -3,7 +3,7 @@
"""
This file contains ops for ViT attention to be compatible with torch.compile
as there are operations here not supported by torch.compile (for instance,
-`to_list` in xformers attn, or `.item()` in flash attention)
+`.item()` in flash attention)
Using these ops and wrapping vision blocks with `torch.compile` can speed up
throughput in vision models by ~5% relative on H100, and improve token
@@ -19,42 +19,6 @@ import torch.nn.functional as F
from vllm.utils.torch_utils import direct_register_custom_op
-def xformers_attn_seqlens_wrapper(
- q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
-) -> torch.Tensor:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
-
- attn_bias = BlockDiagonalMask.from_seqlens(
- q_seqlen=seqlens.tolist(), kv_seqlen=None, device=q.device
- )
- context_layer = xops.memory_efficient_attention_forward(
- q, k, v, attn_bias=attn_bias, p=0, scale=None
- )
- context_layer = einops.rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
- return context_layer
-
-
-def xformers_attn_seqlens_wrapper_fake(
- q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
-) -> torch.Tensor:
- b, s, h, d = q.shape
- return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
-
-
-direct_register_custom_op(
- op_name="xformers_attn_seqlens_wrapper",
- op_func=xformers_attn_seqlens_wrapper,
- fake_impl=xformers_attn_seqlens_wrapper_fake,
-)
-
-
-def vit_xformers_attn_wrapper(
- q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
-) -> torch.Tensor:
- return torch.ops.vllm.xformers_attn_seqlens_wrapper(q, k, v, seqlens)
-
-
def flash_attn_maxseqlen_wrapper(
q: torch.Tensor,
k: torch.Tensor,
diff --git a/vllm/attention/selector.py b/vllm/attention/selector.py
index 1a092db9ce378..ad19b58aa155c 100644
--- a/vllm/attention/selector.py
+++ b/vllm/attention/selector.py
@@ -12,7 +12,11 @@ import torch
import vllm.envs as envs
from vllm.attention.backends.abstract import AttentionBackend
-from vllm.attention.backends.registry import AttentionBackendEnum
+from vllm.attention.backends.registry import (
+ MAMBA_TYPE_TO_BACKEND_MAP,
+ AttentionBackendEnum,
+ MambaAttentionBackendEnum,
+)
from vllm.config.cache import CacheDType
from vllm.logger import init_logger
from vllm.utils import STR_BACKEND_ENV_VAR
@@ -32,7 +36,14 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
* None otherwise
"""
backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
- return None if backend_name is None else AttentionBackendEnum[backend_name]
+ if backend_name is None:
+ return None
+ if backend_name == "XFORMERS":
+ raise ValueError(
+ "Attention backend 'XFORMERS' has been removed (See PR #29262 for "
+ "details). Please select a supported attention backend."
+ )
+ return AttentionBackendEnum[backend_name]
# Global state allows a particular choice of backend
@@ -197,6 +208,33 @@ def _cached_get_attn_backend(
return backend
+def get_mamba_attn_backend(
+ mamba_type: str,
+) -> type[AttentionBackend]:
+ """Select which mamba attention backend to use and lazily import it."""
+ return _cached_get_mamba_attn_backend(mamba_type)
+
+
+@cache
+def _cached_get_mamba_attn_backend(
+ mamba_type: str,
+) -> type[AttentionBackend]:
+ assert mamba_type and isinstance(mamba_type, str)
+
+ selected_backend = None
+ try:
+ backend_name = MAMBA_TYPE_TO_BACKEND_MAP[mamba_type]
+ selected_backend = MambaAttentionBackendEnum[backend_name]
+ except KeyError as e:
+ raise ValueError(
+ f"Invalid mamba attention backend type: '{backend_name}'. Valid "
+ f"backends are: {list(MambaAttentionBackendEnum.__members__.keys())}"
+ ) from e
+
+ mamba_attn_backend = selected_backend.get_class()
+ return mamba_attn_backend
+
+
@contextmanager
def global_force_attn_backend_context_manager(
attn_backend: AttentionBackendEnum,
diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py
index 60ef6eef21663..1e66f21ff6388 100644
--- a/vllm/compilation/backends.py
+++ b/vllm/compilation/backends.py
@@ -4,12 +4,14 @@
import ast
import dataclasses
import hashlib
+import json
import operator
import os
import pprint
import time
from collections.abc import Callable, Sequence
from contextlib import contextmanager
+from functools import partial
from typing import Any
import torch
@@ -23,7 +25,9 @@ from vllm.compilation.partition_rules import (
should_split,
)
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
+from vllm.config.utils import hash_factors
from vllm.logger import init_logger
+from vllm.logging_utils import lazy
from vllm.platforms import current_platform
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -580,35 +584,47 @@ class VllmBackend:
def __call__(
self, graph: fx.GraphModule, example_inputs
) -> VllmSerializableFunction:
- from .caching import _compute_code_hash, compilation_config_hash_factors
-
vllm_config = self.vllm_config
+ # Minimal hashing here with existing utilities, reused below.
+
+ env_factors = envs.compile_factors()
+ env_hash = hash_factors(env_factors)
+ # Compute config/compiler/code hashes once and reuse
+ config_hash = vllm_config.compute_hash()
+ compiler_hash = self.compiler_manager.compute_hash(vllm_config)
+ forward_code_files = list(sorted(self.compilation_config.traced_files))
+
+ logger.debug(
+ "Traced files (to be considered for compilation cache):\n%s",
+ lazy(lambda: "\n".join(forward_code_files)),
+ )
+ hash_content = []
+ for filepath in forward_code_files:
+ hash_content.append(filepath)
+ if filepath == "":
+ # This means the function was dynamically generated, with
+ # e.g. exec(). We can't actually check these.
+ continue
+ try:
+ with open(filepath) as f:
+ hash_content.append(f.read())
+ except Exception:
+ logger.warning("Failed to read file %s", filepath)
+ continue
+ code_hash = hashlib.sha256("\n".join(hash_content).encode()).hexdigest()
+ # Clear after consumption
+ self.compilation_config.traced_files.clear()
if not self.compilation_config.cache_dir:
# no provided cache dir, generate one based on the known factors
# that affects the compilation. if none of the factors change,
# the cache dir will be the same so that we can reuse the compiled
# graph.
-
- factors = compilation_config_hash_factors(vllm_config)
- # 2. factors come from the code files that are traced by Dynamo (
- # it mainly summarizes how the model is used in forward pass)
- code_hash = _compute_code_hash(self.compilation_config.traced_files)
- self.compilation_config.traced_files.clear()
- factors.append(code_hash)
-
- # 3. compiler hash
- compiler_hash = self.compiler_manager.compute_hash(vllm_config)
- factors.append(compiler_hash)
-
- # combine all factors to generate the cache dir
- hash_key = hashlib.md5(
- str(factors).encode(), usedforsecurity=False
- ).hexdigest()[:10]
-
+ factors = [env_hash, config_hash, code_hash, compiler_hash]
+ # Use SHA-256 for cache key hashing to be consistent across
+ # compute_hash functions. Truncate for a short cache dir name.
+ hash_key = hashlib.sha256(str(factors).encode()).hexdigest()[:10]
cache_dir = os.path.join(
- envs.VLLM_CACHE_ROOT,
- "torch_compile_cache",
- hash_key,
+ envs.VLLM_CACHE_ROOT, "torch_compile_cache", hash_key
)
self.compilation_config.cache_dir = cache_dir
@@ -621,6 +637,7 @@ class VllmBackend:
os.makedirs(local_cache_dir, exist_ok=True)
self.compilation_config.local_cache_dir = local_cache_dir
+ # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE.
disable_cache = not is_compile_cache_enabled(
self.compilation_config.inductor_compile_config
)
@@ -638,6 +655,50 @@ class VllmBackend:
local_cache_dir, disable_cache, self.prefix
)
+ # Reuses existing cache key
+
+ logger.debug(
+ "torch.compile cache factors: env=%s cfg=%s comp=%s code=%s dir=%s",
+ env_hash,
+ config_hash,
+ compiler_hash,
+ code_hash,
+ local_cache_dir,
+ )
+
+ # Persist and log only hash-relevant factors together.
+ try:
+ logger.debug(
+ "Compile env factors (raw):\n%s\nVllm config hash: %s",
+ lazy(partial(pprint.pformat, env_factors, width=120)),
+ config_hash,
+ )
+ meta_path = os.path.join(local_cache_dir, "cache_key_factors.json")
+ if not os.path.exists(meta_path):
+ with open(meta_path, "w") as f:
+ json.dump(
+ {
+ "env": env_factors, # raw factors used for env_hash
+ "config_hash": config_hash,
+ "code_hash": code_hash,
+ "compiler_hash": compiler_hash,
+ },
+ f,
+ indent=2,
+ sort_keys=True,
+ )
+ except Exception:
+ # Best-effort only; metadata write failures are non-fatal.
+ logger.warning(
+ (
+ "Could not write compile cache metadata at %s; continuing without "
+ "metadata. Compiled cache remains valid; diagnostics may be "
+ "limited."
+ ),
+ local_cache_dir,
+ exc_info=True,
+ )
+
# when dynamo calls the backend, it means the bytecode
# transform and analysis are done
compilation_counter.num_graphs_seen += 1
diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py
index 16e34c2711e9f..6297d9f995aa4 100644
--- a/vllm/compilation/caching.py
+++ b/vllm/compilation/caching.py
@@ -12,6 +12,7 @@ from torch.utils import _pytree as pytree
import vllm.envs as envs
from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.utils import hash_factors
from vllm.logger import init_logger
try:
@@ -115,7 +116,8 @@ class VllmSerializableFunction(SerializableCallable):
the AOT compiled path.
"""
compile_inputs = [
- inp or example_inputs[i] for i, inp in enumerate(fn.example_inputs)
+ inp if inp is not None else example_inputs[i]
+ for i, inp in enumerate(fn.example_inputs)
]
with tracing(TracingContext(fake_mode)):
fn.optimized_call = vllm_backend(
@@ -138,7 +140,7 @@ def compilation_config_hash_factors(vllm_config: VllmConfig) -> list[str]:
factors = []
# 0. factors come from the env, for example, The values of
# VLLM_PP_LAYER_PARTITION will affect the computation graph.
- env_hash = envs.compute_hash()
+ env_hash = hash_factors(envs.compile_factors())
factors.append(env_hash)
# 1. factors come from the vllm_config (it mainly summarizes how the
diff --git a/vllm/compilation/matcher_utils.py b/vllm/compilation/matcher_utils.py
index 38eb4e5301a18..e4cd063d2aee1 100644
--- a/vllm/compilation/matcher_utils.py
+++ b/vllm/compilation/matcher_utils.py
@@ -162,12 +162,10 @@ class MatcherRMSNorm(MatcherCustomOp):
weight: torch.Tensor,
) -> torch.Tensor:
result = torch.empty_like(input)
- # TODO: support non-contiguous input for RMSNorm and remove this
- input_contiguous = input.contiguous()
_, result = auto_functionalized(
RMS_OP,
result=result,
- input=input_contiguous,
+ input=input,
weight=weight,
epsilon=self.epsilon,
)
diff --git a/vllm/compilation/pass_manager.py b/vllm/compilation/pass_manager.py
index 0e8bb2fc97351..fe2547d7fecaf 100644
--- a/vllm/compilation/pass_manager.py
+++ b/vllm/compilation/pass_manager.py
@@ -127,7 +127,7 @@ class PostGradPassManager(CustomGraphPass):
affects compilation caching. Its uuid depends on the UUIDs of all
dependent passes and the pass config. See InductorPass for more info.
"""
- state = {"pass_config": self.pass_config.uuid(), "passes": []}
+ state = {"pass_config": self.pass_config.compute_hash(), "passes": []}
for pass_ in self.passes:
state["passes"].append(pass_.uuid())
state["passes"].append(self.fix_functionalization.uuid())
diff --git a/vllm/config/cache.py b/vllm/config/cache.py
index 864cf1be81b20..ef6928d8ebd5c 100644
--- a/vllm/config/cache.py
+++ b/vllm/config/cache.py
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import hashlib
from dataclasses import field
from typing import TYPE_CHECKING, Any, Literal
@@ -74,8 +73,8 @@ class CacheConfig:
sliding_window: int | None = None
"""Sliding window size for the KV cache. This is primarily set in
`ModelConfig` and that value should be manually duplicated here."""
- enable_prefix_caching: bool | None = None
- """Whether to enable prefix caching. Enabled by default for V1."""
+ enable_prefix_caching: bool = True
+ """Whether to enable prefix caching."""
prefix_caching_hash_algo: PrefixCachingHashAlgo = "sha256"
"""Set the hash algorithm for prefix caching:\n
- "sha256" uses Pickle for object serialization before hashing.\n
@@ -160,13 +159,29 @@ class CacheConfig:
excluding anything before input ids/embeddings and after
the final hidden states.
"""
- factors: list[Any] = []
- factors.append(self.cache_dtype)
- factors.append(self.mamba_cache_dtype)
- factors.append(self.mamba_ssm_cache_dtype)
- # `cpu_offload_gb` does not use `torch.compile` yet.
- hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
- return hash_str
+ ignored_factors = {
+ # Runtime/derived knobs that don't affect compiled graph shape
+ "gpu_memory_utilization",
+ "swap_space",
+ "is_attention_free",
+ "num_gpu_blocks_override",
+ "enable_prefix_caching",
+ "prefix_caching_hash_algo",
+ # `cpu_offload_gb` does not use `torch.compile` yet.
+ "cpu_offload_gb",
+ "cpu_kvcache_space_bytes",
+ "mamba_page_size_padded",
+ # Post-init/derived counters
+ "num_gpu_blocks",
+ "num_cpu_blocks",
+ # WIP feature toggle not impacting compiled graph shape
+ "kv_sharing_fast_prefill",
+ }
+
+ from vllm.config.utils import get_hash_factors, hash_factors
+
+ factors = get_hash_factors(self, ignored_factors)
+ return hash_factors(factors)
def metrics_info(self):
# convert cache_config to dict(key: str, value: str) for prometheus
diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py
index 088d0b1af757a..9b5309598d0e2 100644
--- a/vllm/config/compilation.py
+++ b/vllm/config/compilation.py
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import enum
-import hashlib
from collections import Counter
from collections.abc import Callable
from dataclasses import asdict, field
@@ -160,7 +159,7 @@ class PassConfig:
current_platform.get_device_capability().to_int(), {}
)
- def uuid(self):
+ def compute_hash(self) -> str:
"""
Produces a hash unique to the pass configuration.
Any new fields that affect compilation should be added to the hash.
@@ -506,28 +505,33 @@ class CompilationConfig:
def compute_hash(self) -> str:
"""
- WARNING: Whenever a new field is added to this config,
- ensure that it is included in the factors list if
- it affects the computation graph.
-
Provide a hash that uniquely identifies all the configs
that affect the structure of the computation
graph from input ids/embeddings to the final hidden states,
excluding anything before input ids/embeddings and after
the final hidden states.
"""
- factors: list[Any] = []
- factors.append(self.mode)
- factors.append(self.backend)
- factors.append(self.custom_ops)
- factors.append(self.splitting_ops)
- factors.append(self.use_inductor)
- factors.append(self.use_inductor_graph_partition)
- factors.append(self.inductor_compile_config)
- factors.append(self.inductor_passes)
- factors.append(self.pass_config.uuid())
- factors.append(self.compile_cache_save_format)
- return hashlib.sha256(str(factors).encode()).hexdigest()
+ # Opt-out: default-include declared fields; keep a tiny exclude set;
+ # normalize types; keep SHA-256. For nested opaque configs, include a
+ # stable identifier (e.g., pass_config.compute_hash()) instead of object id.
+
+ ignored_factors = {
+ # Paths/dirs and runtime/metrics that don’t affect compiled graph
+ "debug_dump_path",
+ "cache_dir",
+ "local_cache_dir",
+ "bs_to_padded_graph_size",
+ "traced_files",
+ "compilation_time",
+ "static_forward_context",
+ "pass_config", # handled separately below
+ }
+
+ from vllm.config.utils import get_hash_factors, hash_factors
+
+ factors = get_hash_factors(self, ignored_factors)
+ factors["pass_config"] = self.pass_config.compute_hash()
+ return hash_factors(factors)
def __repr__(self) -> str:
exclude = {
@@ -660,6 +664,8 @@ class CompilationConfig:
is_torch_equal_or_newer("2.9.0.dev")
and "combo_kernels" not in self.inductor_compile_config
and "benchmark_combo_kernel" not in self.inductor_compile_config
+ # (fixme @boyuan) combo kernel does not support cpu yet.
+ and not current_platform.is_cpu()
):
# use horizontal fusion, which is useful for fusing qk-norm and
# qk-rope when query and key have different shapes.
@@ -917,7 +923,7 @@ class CompilationConfig:
self, uniform_decode_query_len: int, tensor_parallel_size: int
):
multiple_of = uniform_decode_query_len
- if tensor_parallel_size > 1:
+ if tensor_parallel_size > 1 and self.pass_config.enable_sequence_parallelism:
multiple_of = max(uniform_decode_query_len, tensor_parallel_size)
if (
multiple_of % uniform_decode_query_len != 0
@@ -944,14 +950,18 @@ class CompilationConfig:
)
)
+ if len(rounded_sizes) == 0 and multiple_of <= self.max_cudagraph_capture_size:
+ # if one valid but would be round_down use that
+ rounded_sizes = [multiple_of]
+
if len(rounded_sizes) == 0:
- logger.warning(
- "No valid cudagraph sizes after rounding to multiple of "
- " num_speculative_tokens + 1 (%d); please adjust num_speculative_tokens"
- " or max_cudagraph_capture_size (or cudagraph_capture_sizes)",
- multiple_of,
+ raise ValueError(
+ f"No valid cudagraph sizes after rounding to multiple of {multiple_of} "
+ f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)"
+ f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}"
+ f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})"
+ f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})"
)
- return
self.max_cudagraph_capture_size = rounded_sizes[-1]
self.cudagraph_capture_sizes = rounded_sizes
diff --git a/vllm/config/lora.py b/vllm/config/lora.py
index 84e92eef40077..072e0ec2104f5 100644
--- a/vllm/config/lora.py
+++ b/vllm/config/lora.py
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import hashlib
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, Literal
import torch
from pydantic import ConfigDict, Field, model_validator
@@ -11,7 +11,6 @@ from typing_extensions import Self
from vllm.config.utils import config
from vllm.logger import init_logger
-from vllm.platforms import current_platform
if TYPE_CHECKING:
from vllm.config import ModelConfig
@@ -46,19 +45,6 @@ class LoRAConfig:
`max_loras`."""
lora_dtype: torch.dtype | LoRADType = "auto"
"""Data type for LoRA. If auto, will default to base model dtype."""
- lora_extra_vocab_size: LoRAExtraVocabSize = Field(
- default=256,
- deprecated=(
- "`lora_extra_vocab_size` is deprecated and will be removed "
- "in v0.12.0. Additional vocabulary support for "
- "LoRA adapters is being phased out."
- ),
- )
- """(Deprecated) Maximum size of extra vocabulary that can be present in a
- LoRA adapter. Will be removed in v0.12.0."""
- lora_vocab_padding_size: ClassVar[int] = (
- current_platform.get_lora_vocab_padding_size()
- )
default_mm_loras: dict[str, str] | None = None
"""Dictionary mapping specific modalities to LoRA model paths; this field
is only applicable to multimodal models and should be leveraged when a
@@ -87,8 +73,6 @@ class LoRAConfig:
factors.append(self.max_loras)
factors.append(self.fully_sharded_loras)
factors.append(self.lora_dtype)
- factors.append(self.lora_extra_vocab_size)
- factors.append(self.lora_vocab_padding_size)
hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
return hash_str
diff --git a/vllm/config/model.py b/vllm/config/model.py
index b3a28af6de389..49688e17cf932 100644
--- a/vllm/config/model.py
+++ b/vllm/config/model.py
@@ -1,8 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import hashlib
-import json
import warnings
from collections.abc import Callable
from dataclasses import InitVar, field
@@ -13,12 +11,13 @@ import torch
from pydantic import ConfigDict, SkipValidation, field_validator, model_validator
from pydantic.dataclasses import dataclass
from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
+from transformers.configuration_utils import ALLOWED_LAYER_TYPES
import vllm.envs as envs
from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig
from vllm.config.pooler import PoolerConfig
from vllm.config.scheduler import RunnerType
-from vllm.config.utils import assert_hashable, config, getattr_iter
+from vllm.config.utils import config, getattr_iter
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.transformers_utils.config import (
@@ -35,8 +34,11 @@ from vllm.transformers_utils.config import (
try_get_tokenizer_config,
uses_mrope,
)
+from vllm.transformers_utils.gguf_utils import (
+ maybe_patch_hf_config_from_gguf,
+)
from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
-from vllm.transformers_utils.utils import maybe_model_redirect
+from vllm.transformers_utils.utils import check_gguf_file, maybe_model_redirect
from vllm.utils.import_utils import LazyLoader
from vllm.utils.torch_utils import common_broadcastable_dtype
@@ -79,7 +81,7 @@ TaskOption = Literal[
"transcription",
"draft",
]
-TokenizerMode = Literal["auto", "slow", "mistral", "custom"]
+TokenizerMode = Literal["auto", "hf", "slow", "mistral", "custom"]
ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
LogprobsMode = Literal[
"raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs"
@@ -128,7 +130,8 @@ class ModelConfig:
name or path will be used."""
tokenizer_mode: TokenizerMode = "auto"
"""Tokenizer mode:\n
- - "auto" will use the fast tokenizer if available.\n
+ - "auto" will use "hf" tokenizer if Mistral's tokenizer is not available.\n
+ - "hf" will use the fast tokenizer if available.\n
- "slow" will always use the slow tokenizer.\n
- "mistral" will always use the tokenizer from `mistral_common`.\n
- "custom" will use --tokenizer to select the preregistered tokenizer."""
@@ -144,9 +147,12 @@ class ModelConfig:
- "bfloat16" for a balance between precision and range.\n
- "float" is shorthand for FP32 precision.\n
- "float32" for FP32 precision."""
- seed: int | None = None
- """Random seed for reproducibility. Initialized to None in V0, but
- initialized to 0 in V1."""
+ seed: int = 0
+ """Random seed for reproducibility.
+
+ We must set the global seed because otherwise,
+ different tensor parallel workers would sample different tokens,
+ leading to inconsistent results."""
hf_config: PretrainedConfig = field(init=False)
"""The Hugging Face config of the model."""
hf_text_config: PretrainedConfig = field(init=False)
@@ -236,8 +242,8 @@ class ModelConfig:
first one."""
config_format: str | ConfigFormat = "auto"
"""The format of the model config to load:\n
- - "auto" will try to load the config in hf format if available else it
- will try to load in mistral format.\n
+ - "auto" will try to load the config in hf format if available after trying
+ to load in mistral format.\n
- "hf" will load the config in hf format.\n
- "mistral" will load the config in mistral format."""
hf_token: bool | str | None = None
@@ -320,50 +326,50 @@ class ModelConfig:
excluding anything before input ids/embeddings and after
the final hidden states.
"""
- factors: list[Any] = []
- factors.append(self.model)
- factors.append(self.dtype)
- factors.append(self.quantization)
- factors.append(self.revision)
- factors.append(self.code_revision)
- factors.append(self.max_model_len)
- factors.append(self.max_logprobs)
- factors.append(self.disable_sliding_window)
- factors.append(self.trust_remote_code)
- factors.append(self.generation_config)
- factors.append(self.model_impl)
- factors.append(self.override_generation_config)
- factors.append(self.video_pruning_rate)
- factors.append(self.enable_prompt_embeds)
+ ignored_factors = {
+ "runner",
+ "convert",
+ "task",
+ "tokenizer",
+ "tokenizer_mode",
+ "seed",
+ "hf_config_path",
+ "allowed_local_media_path",
+ "allowed_media_domains",
+ "tokenizer_revision",
+ "spec_target_max_model_len",
+ "enforce_eager",
+ "logprobs_mode",
+ "disable_cascade_attn",
+ "skip_tokenizer_init",
+ "enable_prompt_embeds",
+ "served_model_name",
+ "config_format",
+ "hf_token",
+ "hf_overrides",
+ "logits_processor_pattern",
+ "enable_sleep_mode",
+ "override_attention_dtype",
+ "logits_processors",
+ "io_processor_plugin",
+ "pooler_config",
+ "override_pooler_config",
+ "multimodal_config",
+ "limit_mm_per_prompt",
+ "media_io_kwargs",
+ "mm_processor_kwargs",
+ "mm_processor_cache_gb",
+ "mm_processor_cache_type",
+ "mm_shm_cache_max_object_size_mb",
+ "mm_encoder_tp_mode",
+ "interleave_mm_strings",
+ "skip_mm_profiling",
+ }
- # hf_config can control how the model looks!
- try:
- hf_config_json = self.hf_config.to_json_string(use_diff=False)
- except TypeError:
- from transformers import PretrainedConfig
+ from vllm.config.utils import get_hash_factors, hash_factors
- from vllm.utils.jsontree import json_map_leaves
-
- # Handle nested HF configs with unserializable values gracefully
- hf_config_json = (
- json.dumps(
- json_map_leaves(
- lambda v: v.to_dict()
- if isinstance(v, PretrainedConfig)
- else str(v),
- self.hf_config.to_dict(),
- ),
- indent=2,
- sort_keys=True,
- )
- + "\n"
- )
-
- factors.append(hf_config_json)
-
- str_factors = str(factors)
- assert_hashable(str_factors)
- return hashlib.sha256(str(factors).encode()).hexdigest()
+ factors = get_hash_factors(self, ignored_factors)
+ return hash_factors(factors)
def _update_nested(
self,
@@ -413,7 +419,7 @@ class ModelConfig:
def __post_init__(
self,
# Multimodal config init vars
- limit_mm_per_prompt: dict[str, int] | None,
+ limit_mm_per_prompt: dict[str, int | dict[str, int]] | None,
enable_mm_embeds: bool | None,
media_io_kwargs: dict[str, dict[str, Any]] | None,
mm_processor_kwargs: dict[str, Any] | None,
@@ -426,23 +432,6 @@ class ModelConfig:
skip_mm_profiling: bool | None,
video_pruning_rate: float | None,
) -> None:
- # Set the default seed to 0 in V1.
- # NOTE(woosuk): In V1, we use separate processes for workers (unless
- # VLLM_ENABLE_V1_MULTIPROCESSING=0), so setting a seed here
- # doesn't affect the user process. However, without a consistent seed,
- # different tensor parallel workers would sample different tokens,
- # leading to inconsistent results.
- if self.seed is None:
- self.seed = 0
- if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
- logger.warning(
- "The global random seed is set to %d. Since "
- "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
- "affect the random state of the Python process that "
- "launched vLLM.",
- self.seed,
- )
-
# Keep set served_model_name before maybe_model_redirect(self.model)
self.served_model_name = get_served_model_name(
self.model, self.served_model_name
@@ -450,6 +439,12 @@ class ModelConfig:
self.model = maybe_model_redirect(self.model)
# The tokenizer is consistent with the model by default.
if self.tokenizer is None:
+ if check_gguf_file(self.model):
+ raise ValueError(
+ "Using a tokenizer is mandatory when loading a GGUF model. "
+ "Please specify the tokenizer path or name using the "
+ "--tokenizer argument."
+ )
self.tokenizer = self.model
if self.tokenizer_revision is None:
self.tokenizer_revision = self.revision
@@ -508,6 +503,10 @@ class ModelConfig:
hf_overrides_kw=hf_overrides_kw,
hf_overrides_fn=hf_overrides_fn,
)
+ hf_config = maybe_patch_hf_config_from_gguf(
+ self.model,
+ hf_config,
+ )
self.hf_config = hf_config
if dict_overrides:
@@ -1006,6 +1005,8 @@ class ModelConfig:
# Ensure heavy backends are probed last to avoid unnecessary
# imports during override detection (e.g., MXFP4 imports Triton)
"mxfp4",
+ "cpu_gptq",
+ "cpu_awq",
]
quantization_methods = [
q for q in supported_quantization if q not in overrides
@@ -1137,12 +1138,6 @@ class ModelConfig:
self,
parallel_config: ParallelConfig,
) -> None:
- if parallel_config.distributed_executor_backend == "external_launcher":
- assert self.seed is not None, (
- "Seed must be set when using external launcher backend to "
- "make sure sampling results are the same across workers."
- )
-
total_num_attention_heads = getattr(
self.hf_text_config, "num_attention_heads", 0
)
@@ -1353,11 +1348,7 @@ class ModelConfig:
# Coerce to 0 if explicitly set to None
return num_experts or 0
- def get_layers_start_end_indices(
- self, parallel_config: ParallelConfig
- ) -> tuple[int, int]:
- from vllm.distributed.utils import get_pp_indices
-
+ def get_total_num_hidden_layers(self) -> int:
if (
self.hf_text_config.model_type == "deepseek_mtp"
or self.hf_config.model_type == "mimo_mtp"
@@ -1377,6 +1368,15 @@ class ModelConfig:
total_num_hidden_layers = getattr(
self.hf_text_config, "num_hidden_layers", 0
)
+ return total_num_hidden_layers
+
+ def get_layers_start_end_indices(
+ self, parallel_config: ParallelConfig
+ ) -> tuple[int, int]:
+ from vllm.distributed.utils import get_pp_indices
+
+ total_num_hidden_layers = self.get_total_num_hidden_layers()
+
# the layout order is: DP x PP x TP
pp_rank = (
parallel_config.rank // parallel_config.tensor_parallel_size
@@ -2077,31 +2077,32 @@ def _get_and_verify_max_len(
)
derived_max_model_len = default_max_len
- rope_scaling = getattr(hf_config, "rope_scaling", None)
+ # In Transformers v5 rope_parameters could be TypedDict or dict[str, TypedDict].
+ # To simplify the verification, we convert it to dict[str, TypedDict].
+ rope_parameters = getattr(hf_config, "rope_parameters", None)
+ if rope_parameters and not set(rope_parameters.keys()).issubset(
+ ALLOWED_LAYER_TYPES
+ ):
+ rope_parameters = {"": rope_parameters}
+
# NOTE(woosuk): Gemma3's max_model_len (128K) is already scaled by RoPE
# scaling, so we skip applying the scaling factor again.
- if rope_scaling is not None and "gemma3" not in hf_config.model_type:
- # No need to consider "type" key because of patch_rope_scaling when
- # loading HF config
- rope_type = rope_scaling["rope_type"]
+ if rope_parameters is not None and "gemma3" not in hf_config.model_type:
+ scaling_factor = 1.0
+ for rp in rope_parameters.values():
+ # No need to consider "type" key because of patch_rope_parameters when
+ # loading HF config
+ rope_type = rp["rope_type"]
- if rope_type not in ("su", "longrope", "llama3"):
- if disable_sliding_window:
- # TODO(robertgshaw): Find a model that supports rope_scaling
- # with sliding window to see if this case should be allowed.
- raise NotImplementedError(
- "Disabling sliding window is not supported for models "
- "with rope_scaling. Please raise an issue so we can "
- "investigate."
- )
+ if rope_type not in ("su", "longrope", "llama3"):
+ # NOTE: rope_type == "default" does not define factor https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/modeling_rope_utils.py
+ # NOTE: This assumes all layer types have the same scaling factor.
+ scaling_factor = rp.get("factor", scaling_factor)
- # NOTE: rope_type == "default" does not define factor
- # https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/modeling_rope_utils.py
- scaling_factor = rope_scaling.get("factor", 1.0)
-
- if rope_type == "yarn":
- derived_max_model_len = rope_scaling["original_max_position_embeddings"]
- derived_max_model_len *= scaling_factor
+ if rope_type == "yarn":
+ derived_max_model_len = rp["original_max_position_embeddings"]
+ # Do this outside loop since all layer types should have the same scaling
+ derived_max_model_len *= scaling_factor
if encoder_config and "max_seq_length" in encoder_config:
derived_max_model_len = encoder_config["max_seq_length"]
@@ -2111,7 +2112,9 @@ def _get_and_verify_max_len(
if max_model_len is None:
# For LongRoPE, default to original_max_position_embeddings to avoid
# performance degradation for shorter sequences
- if rope_scaling is not None and rope_scaling["rope_type"] == "longrope":
+ if rope_parameters is not None and any(
+ rp["rope_type"] == "longrope" for rp in rope_parameters.values()
+ ):
max_model_len = int(
getattr(
hf_config, "original_max_position_embeddings", derived_max_model_len
@@ -2128,16 +2131,7 @@ def _get_and_verify_max_len(
# that will be bigger than derived_max_model_len. We compare user input
# with model_max_length and allow this override when it's smaller.
model_max_length = getattr(hf_config, "model_max_length", None)
- if model_max_length is not None and max_model_len <= model_max_length:
- if disable_sliding_window:
- # TODO(robertgshaw): Find a model that has model_max_length
- # with sliding window to see if this case should be allowed.
- raise NotImplementedError(
- "Disabling sliding window is not supported for models "
- "model_max_length in the config. Please raise an issue "
- "so we can investigate."
- )
- else:
+ if model_max_length is None or max_model_len > model_max_length:
msg = (
f"User-specified max_model_len ({max_model_len}) is greater "
f"than the derived max_model_len ({max_len_key}="
diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py
index 9f62b35ed515c..00a81a319bf72 100644
--- a/vllm/config/multimodal.py
+++ b/vllm/config/multimodal.py
@@ -173,6 +173,12 @@ class MultiModalConfig:
# We need to import the real type here (deferred to avoid circular import).
from vllm.attention.backends.registry import AttentionBackendEnum
+ if isinstance(value, str) and value.upper() == "XFORMERS":
+ raise ValueError(
+ "Attention backend 'XFORMERS' has been removed (See PR #29262 for "
+ "details). Please select a supported attention backend."
+ )
+
if value is None or isinstance(value, AttentionBackendEnum):
return value
diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
index 9a6326d62e82e..4b0236d8de3f5 100644
--- a/vllm/config/parallel.py
+++ b/vllm/config/parallel.py
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import hashlib
import os
from typing import TYPE_CHECKING, Any, Literal
@@ -72,6 +71,8 @@ class ParallelConfig:
"""Number of pipeline parallel groups."""
tensor_parallel_size: int = 1
"""Number of tensor parallel groups."""
+ prefill_context_parallel_size: int = 1
+ """Number of prefill context parallel groups."""
data_parallel_size: int = 1
"""Number of data parallel groups. MoE layers will be sharded according to
the product of the tensor parallel size and data parallel size."""
@@ -240,14 +241,25 @@ class ParallelConfig:
needs to be divisible by dcp_size."""
dcp_kv_cache_interleave_size: int = 1
- """Interleave size of kv_cache storage while using dcp or cp > 1,
- store interleave_size tokens on (d)cp i,
- then store next interleave_size tokens on (d)cp i+1.
- Interleave_size=1: token-level align, token i is stored on rank i % (d)cp_size.
- Interleave_size=block_size: block-level align, first fill the block on first rank,
- token is stored on rank i+1 block j after rank i block j is full.
- Block_size should be greater than or equal to dcp_kv_cache_interleave_size.
- Block_size should be divisible by dcp_kv_cache_interleave_size.
+ """
+ Interleave size of kv_cache storage while using DCP.
+ dcp_kv_cache_interleave_size has been replaced by cp_kv_cache_interleave_size,
+ and will be deprecated when PCP is fully supported.
+
+ """
+ cp_kv_cache_interleave_size: int = 1
+ """Interleave size of kv_cache storage while using DCP or PCP.
+ For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
+ and `total_cp_world_size = pcp_world_size * dcp_world_szie`.
+ store interleave_size tokens on total_cp_rank i,
+ then store next interleave_size tokens on taotal_cp_rank i+1.
+ Interleave_size=1: token-level alignment, where token `i` is stored on
+ total_cp_rank `i % total_cp_world_size`.
+ Interleave_size=block_size: block-level alignment, where tokens are
+ first populated to the preceding ranks. Tokens are then stored
+ in (rank i+1, block j) only after (rank i, block j) is fully occupied.
+ Block_size should be greater than or equal to cp_kv_cache_interleave_size.
+ Block_size should be divisible by cp_kv_cache_interleave_size.
"""
_api_process_count: int = Field(default=1, gt=0)
@@ -312,6 +324,11 @@ class ParallelConfig:
"num_redundant_experts."
)
+ if self.prefill_context_parallel_size > 1:
+ raise ValueError(
+ "Prefill context parallelism is not fully supported. "
+ "Please set prefill_context_parallel_size to 1."
+ )
return self
@property
@@ -448,19 +465,41 @@ class ParallelConfig:
This hash is also used for DP worker configuration validation
to prevent hangs from mismatched collective communication patterns.
"""
- factors: list[Any] = []
- factors.append(self.pipeline_parallel_size)
- factors.append(self.tensor_parallel_size)
- factors.append(self.enable_expert_parallel)
- factors.append(self.data_parallel_size)
- factors.append(self.all2all_backend)
- factors.append(self.enable_eplb)
- if self.enable_eplb:
- factors.append(self.eplb_config.log_balancedness)
- factors.append(self.eplb_config.window_size)
- factors.append(self.eplb_config.step_interval)
- factors.append(self.eplb_config.num_redundant_experts)
- return hashlib.sha256(str(factors).encode()).hexdigest()
+ ignored_factors = {
+ # Derived/runtime topology, networking, or launch details
+ "data_parallel_rank",
+ "data_parallel_rank_local",
+ "data_parallel_backend",
+ "data_parallel_external_lb",
+ "data_parallel_hybrid_lb",
+ "data_parallel_master_ip",
+ "data_parallel_master_port",
+ "_data_parallel_master_port_list",
+ "data_parallel_rpc_port",
+ "rank",
+ "master_addr",
+ "master_port",
+ "node_rank",
+ "nnodes",
+ "max_parallel_loading_workers",
+ "disable_custom_all_reduce",
+ "ray_workers_use_nsight",
+ "ray_runtime_env",
+ "placement_group",
+ "distributed_executor_backend",
+ "worker_cls",
+ "sd_worker_cls",
+ "worker_extension_cls",
+ "_api_process_count",
+ "_api_process_rank",
+ }
+
+ from vllm.config.utils import get_hash_factors, hash_factors
+
+ factors = get_hash_factors(self, ignored_factors)
+ # Explicitly include backend affecting env factor as before
+ factors["VLLM_ALL2ALL_BACKEND"] = str(envs.VLLM_ALL2ALL_BACKEND)
+ return hash_factors(factors)
def __post_init__(self) -> None:
# Set all2all_backend from env var if not specified, with deprecation warning
@@ -508,7 +547,11 @@ class ParallelConfig:
)
# Continue with the rest of the initialization
- self.world_size = self.pipeline_parallel_size * self.tensor_parallel_size
+ self.world_size = (
+ self.pipeline_parallel_size
+ * self.tensor_parallel_size
+ * self.prefill_context_parallel_size
+ )
if self.distributed_executor_backend == "external_launcher":
logger.info("Using external launcher for distributed inference.")
diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py
index 8194295ffedb6..b6078706daacf 100644
--- a/vllm/config/scheduler.py
+++ b/vllm/config/scheduler.py
@@ -62,15 +62,6 @@ class SchedulerConfig:
"""For chunked prefill, a request is considered long if the prompt is
longer than this number of tokens."""
- num_lookahead_slots: int = Field(default=0, ge=0)
- """The number of slots to allocate per sequence per
- step, beyond the known token ids. This is used in speculative
- decoding to store KV activations of tokens which may or may not be
- accepted.
-
- NOTE: This will be replaced by speculative config in the future; it is
- present to enable correctness tests until then."""
-
enable_chunked_prefill: bool = True
"""If True, prefill requests can be chunked based
on the remaining `max_num_batched_tokens`.
diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index 13a8632413d91..d7c019c73d598 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -9,6 +9,7 @@ from pydantic import Field, SkipValidation, model_validator
from pydantic.dataclasses import dataclass
from typing_extensions import Self
+from vllm.config.model import ModelConfig
from vllm.config.parallel import ParallelConfig
from vllm.config.utils import config
from vllm.logger import init_logger
@@ -18,10 +19,8 @@ if TYPE_CHECKING:
from transformers import PretrainedConfig
import vllm.model_executor.layers.quantization as me_quant
- from vllm.config import ModelConfig
else:
PretrainedConfig = Any
- ModelConfig = Any
me_quant = LazyLoader(
"model_executor", globals(), "vllm.model_executor.layers.quantization"
@@ -316,10 +315,6 @@ class SpeculativeConfig:
self.prompt_lookup_min = 0
if self.model is not None:
- # TODO: Move this import to the top once `ModelConfig`
- # lives in `vllm.config.model`.
- from vllm.config import ModelConfig
-
self.draft_model_config = ModelConfig(
model=self.model,
runner="draft",
@@ -634,16 +629,6 @@ class SpeculativeConfig:
return self
- @property
- def num_lookahead_slots(self) -> int:
- """The number of additional slots the scheduler should allocate per
- step, in addition to the slots allocated for each known token.
-
- This is equal to the number of speculative tokens, as each speculative
- token must be scored.
- """
- return self.num_speculative_tokens
-
def use_eagle(self) -> bool:
return self.method in ("eagle", "eagle3", "mtp")
diff --git a/vllm/config/utils.py b/vllm/config/utils.py
index 7e0878d96bbd6..02f2b75f608f1 100644
--- a/vllm/config/utils.py
+++ b/vllm/config/utils.py
@@ -3,14 +3,19 @@
"""Utility functions for vLLM config dataclasses."""
import ast
+import enum
+import hashlib
import inspect
+import json
+import pathlib
import textwrap
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping, Sequence, Set
from dataclasses import MISSING, Field, field, fields, is_dataclass, replace
from itertools import pairwise
from typing import TYPE_CHECKING, Any, Protocol, TypeVar
import regex as re
+import torch
from pydantic.fields import FieldInfo
from typing_extensions import runtime_checkable
@@ -176,3 +181,115 @@ def update_config(config: ConfigT, overrides: dict[str, Any]) -> ConfigT:
)
processed_overrides[field_name] = value
return replace(config, **processed_overrides)
+
+
+def normalize_value(x):
+ """Return a stable, JSON-serializable canonical form for hashing.
+ Order: primitives, special types (Enum, callable, torch.dtype, Path), then
+ generic containers (Mapping/Set/Sequence) with recursion.
+ """
+ # Fast path
+ if x is None or isinstance(x, (bool, int, float, str)):
+ return x
+
+ # Enums: tag with FQN to avoid primitive collisions.
+ # Ex: Enum(1) vs int(1) -> ("module.QualName", value).
+ if isinstance(x, enum.Enum):
+ enum_type = f"{x.__class__.__module__}.{x.__class__.__qualname__}"
+ return (enum_type, normalize_value(x.value))
+
+ # Classes (types) are accepted and canonicalized by their fully-qualified
+ # name (module.qualname) for a stable identifier.
+ # Instances are only accepted if they expose uuid(); otherwise they are
+ # rejected to avoid under-hashing object state.
+
+ # Callables: accept classes only; reject funcs/lambdas/methods.
+ # Used by LogitsProcessor types and ModelConfig.hf_overrides.
+ if isinstance(x, type):
+ module = getattr(x, "__module__", "")
+ qual = getattr(x, "__qualname__", getattr(x, "__name__", ""))
+ return ".".join([p for p in (module, qual) if p]) or repr(x)
+
+ # Prefer stable uuid identifiers for objects that provide them, even if
+ # they are callable instances (e.g., InductorPass wrappers).
+ if hasattr(x, "uuid") and callable(getattr(x, "uuid", None)):
+ return x.uuid()
+
+ if callable(x):
+ raise TypeError("normalize_value: function or callable instance unsupported")
+
+ # Torch dtype: stringify (torch.float64 -> "torch.float64").
+ # We rely on the string form here; dtype-bearing fields that need additional
+ # disambiguation should encode that at the config layer.
+ if isinstance(x, torch.dtype):
+ return str(x)
+
+ # Bytes
+ if isinstance(x, (bytes, bytearray)):
+ return x.hex()
+
+ # Paths (canonicalize)
+ if isinstance(x, pathlib.Path):
+ try:
+ return str(x.expanduser().resolve())
+ except Exception:
+ return str(x)
+
+ # Dataclasses: represent as (FQN, sorted(field,value) tuple) for stability.
+ if is_dataclass(x):
+ type_fqn = f"{x.__class__.__module__}.{x.__class__.__qualname__}"
+ items = tuple(
+ (f.name, normalize_value(getattr(x, f.name)))
+ for f in sorted(fields(x), key=lambda f: f.name)
+ )
+ return (type_fqn, items)
+
+ # Containers (generic)
+ if isinstance(x, Mapping):
+ return tuple(sorted((str(k), normalize_value(v)) for k, v in x.items()))
+ if isinstance(x, Set):
+ return tuple(sorted(repr(normalize_value(v)) for v in x))
+ if isinstance(x, Sequence) and not isinstance(x, (str, bytes, bytearray)):
+ return tuple(normalize_value(v) for v in x)
+
+ # PretrainedConfig
+ if hasattr(x, "to_json_string") and callable(x.to_json_string):
+ return x.to_json_string()
+
+ # Unsupported type: e.g., modules, generators, open files, or objects
+ # without a stable JSON/UUID representation. Hard-error to avoid
+ # under-hashing.
+ # If you hit this, either reshape your config to use supported primitives
+ # and containers, or extend normalize_value to provide a stable encoding
+ # (e.g., via uuid() or to_json_string()) for this type.
+ raise TypeError(
+ f"normalize_value: unsupported type '{type(x).__name__}'. "
+ "Ensure config values use supported primitives/containers or add a "
+ "stable representation for this type."
+ )
+
+
+def get_hash_factors(config: ConfigT, ignored_factors: set[str]) -> dict[str, object]:
+ """Gets the factors used for hashing a config class.
+ - Includes all dataclass fields not in `ignored_factors`.
+ - Errors on non-normalizable values.
+ """
+ factors: dict[str, object] = {}
+ for dc_field in fields(config):
+ factor = dc_field.name
+ if factor in ignored_factors:
+ continue
+ value = getattr(config, factor, None)
+ try:
+ factors[factor] = normalize_value(value)
+ except TypeError as e:
+ raise TypeError(
+ f"get_hash_factors: unsupported type for key '{factor}' "
+ f"({type(value).__name__})"
+ ) from e
+ return factors
+
+
+def hash_factors(items: dict[str, object]) -> str:
+ """Return a SHA-256 hex digest of the canonical items structure."""
+ return hashlib.sha256(json.dumps(items, sort_keys=True).encode()).hexdigest()
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index 672b004c4aa56..d64e315b4fe39 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -481,6 +481,14 @@ class VllmConfig:
"Overriding cudagraph_mode to PIECEWISE."
)
self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
+ # prefill context parallel do not support full cudagraphs
+ elif self.parallel_config.prefill_context_parallel_size > 1:
+ logger.warning_once(
+ "Prefill context parallel (PCP) is enabled, which is "
+ "incompatible with full CUDA graphs. "
+ "Overriding cudagraph_mode to PIECEWISE."
+ )
+ self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
elif self.model_config is not None:
if self.model_config.pooler_config is not None:
logger.warning_once(
@@ -610,22 +618,34 @@ class VllmConfig:
# If DCP, ensure the block size is right.
if self.parallel_config.decode_context_parallel_size > 1:
+ if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
+ self.parallel_config.cp_kv_cache_interleave_size
+ != self.parallel_config.dcp_kv_cache_interleave_size
+ ):
+ self.parallel_config.cp_kv_cache_interleave_size = (
+ self.parallel_config.dcp_kv_cache_interleave_size
+ )
+ logger.warning_once(
+ "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
+ "_interleave_size. And dcp-kv-cache-interleave-size will be "
+ "deprecated when PCP is fully supported."
+ )
assert (
- self.parallel_config.dcp_kv_cache_interleave_size
+ self.parallel_config.cp_kv_cache_interleave_size
<= self.cache_config.block_size
and self.cache_config.block_size
- % self.parallel_config.dcp_kv_cache_interleave_size
+ % self.parallel_config.cp_kv_cache_interleave_size
== 0
), (
f"Block_size({self.cache_config.block_size}) should be greater "
- "than or equal to and divisible by dcp_kv_cache_interleave_size "
- f"({self.parallel_config.dcp_kv_cache_interleave_size})."
+ "than or equal to and divisible by cp_kv_cache_interleave_size "
+ f"({self.parallel_config.cp_kv_cache_interleave_size})."
)
assert (
- self.parallel_config.dcp_kv_cache_interleave_size == 1
+ self.parallel_config.cp_kv_cache_interleave_size == 1
or self.speculative_config is None
- ), "MTP with dcp_kv_cache_interleave_size > 1 is not supported now."
+ ), "MTP with cp_kv_cache_interleave_size > 1 is not supported now."
# Do this after all the updates to compilation_config.mode
if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
index f85eb414b2222..74f09278b7bb1 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
@@ -38,7 +38,7 @@ The class provides the following primitives:
import enum
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterable
-from typing import TYPE_CHECKING, Any, Literal, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional
import torch
@@ -47,7 +47,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import KVConnectorOutput
if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionMetadata
+ from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
from vllm.config import VllmConfig
from vllm.distributed.kv_events import KVCacheEvent
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
@@ -142,6 +142,18 @@ class KVConnectorMetadata(ABC): # noqa: B024
class KVConnectorBase_V1(ABC):
+ """
+ Base class for KV connectors.
+
+ Attributes:
+ prefer_cross_layer_blocks (bool): Indicates whether this connector
+ prefers KV blocks that hold KV data for all layers (for speeding
+ up KV data transfers).
+ Defaults to False.
+ """
+
+ prefer_cross_layer_blocks: ClassVar[bool] = False
+
def __init__(
self,
vllm_config: "VllmConfig",
@@ -226,6 +238,23 @@ class KVConnectorBase_V1(ABC):
"""
return
+ def register_cross_layers_kv_cache(
+ self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"]
+ ):
+ """
+ Initialize with a single KV cache tensor used by all layers.
+ The first dimension should be num_layers.
+ This function will only be called for models with uniform layers,
+ and only if the prefers_cross_layer_blocks is set to True.
+ Only one of the functions
+ {register_kv_caches, register_cross_layers_kv_cache} will be called.
+
+ Args:
+ kv_cache: a cross-layers kv cache tensor
+ attn_backend: The attention backend that corresponds to all layers
+ """
+ return
+
def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
"""
Set the xPU-specific ops for copying KV between host and device.
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
index ab2eeed9f6b8a..6acfb73997f25 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
@@ -310,7 +310,6 @@ class LMCacheMPWorkerAdapter:
request_id,
result,
)
- logger.info("Retrieve request for request_id=%s finished", request_id)
# Remove the finished requests from the tracking dicts
for request_id in finished_stores:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
index 55831dc56c803..d1d3e475cc889 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, cast
import torch
import zmq
+from lmcache.integration.vllm.utils import mla_enabled
from lmcache.utils import init_logger as lmcache_init_logger
from vllm.config import VllmConfig
@@ -60,17 +61,44 @@ def reformat_block_ids(block_ids: tuple[list[int], ...] | None) -> list[int]:
return block_ids[0]
+def extract_world_size_and_kv_rank(
+ world_size: int,
+ rank: int,
+ vllm_config: VllmConfig,
+) -> tuple[int, int]:
+ """
+ Convert the rank for the MLA.
+ """
+ use_mla = mla_enabled(vllm_config.model_config)
+ if not use_mla:
+ return world_size, rank
+ else:
+ # Tensor parallel does not change the KV caches for MLA models.
+ # So we need to "exclude" the effect of TP on rank and world size
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
+ # vLLM constructs TP groups first, and then construct other
+ # parallel groups on top of TP groups.
+ # for example, TP=4, PP=2,
+ # TP group: [0, 1, 2, 3], [4, 5, 6, 7]
+ # PP group: [0, 4], [1, 5], [2, 6], [3, 7]
+ # So we can "exclude" the effect of TP by rank // tp_size.
+ return world_size // tp_size, rank // tp_size
+
+
def create_scheduler_adapter(
server_url: str, zmq_context: zmq.Context, vllm_config: VllmConfig
) -> LMCacheMPSchedulerAdapter:
- # TODO: have a helper function to calculate the correct rank and
- # world size for the MLA and other models
+ world_size, kv_rank = extract_world_size_and_kv_rank(
+ vllm_config.parallel_config.world_size,
+ vllm_config.parallel_config.rank,
+ vllm_config,
+ )
return LMCacheMPSchedulerAdapter(
server_url,
zmq_context,
vllm_config.model_config.model,
- vllm_config.parallel_config.world_size,
- vllm_config.parallel_config.rank,
+ world_size,
+ kv_rank,
vllm_config.cache_config.block_size,
)
@@ -78,14 +106,17 @@ def create_scheduler_adapter(
def create_worker_adapter(
server_url: str, zmq_context: zmq.Context, vllm_config: VllmConfig
) -> LMCacheMPWorkerAdapter:
- # TODO: have a helper function to calculate the correct rank and
- # world size for the MLA and other models
+ world_size, kv_rank = extract_world_size_and_kv_rank(
+ vllm_config.parallel_config.world_size,
+ vllm_config.parallel_config.rank,
+ vllm_config,
+ )
return LMCacheMPWorkerAdapter(
server_url,
zmq_context,
vllm_config.model_config.model,
- vllm_config.parallel_config.world_size,
- vllm_config.parallel_config.rank,
+ world_size,
+ kv_rank,
vllm_config.cache_config.block_size,
)
@@ -438,9 +469,6 @@ class LMCacheMPConnector(KVConnectorBase_V1):
ops.append(meta.op)
if len(request_ids) > 0:
- logger.info(
- "HERE! SUBMITTING THE BATCHED RETRIEVE REQUESTS %s", request_ids
- )
self.worker_adapter.batched_submit_retrieve_requests(
request_ids, ops, event
)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
index a70c98b637131..7c0911240493c 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
@@ -677,12 +677,13 @@ class NixlConnectorWorker:
mapping between local and remote TP workers.
"""
- tp_size: int
tp_rank: int
remote_tp_size: dict[EngineId, int]
is_mla: bool
total_num_kv_heads: int
attn_backend: type[AttentionBackend]
+ engine_id: EngineId
+ remote_block_size: dict[EngineId, int]
def __post_init__(self):
# Figure out whether the first dimension of the cache is K/V
@@ -710,8 +711,13 @@ class NixlConnectorWorker:
self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first
)
- block_size: int
- remote_block_size: dict[EngineId, int]
+ @property
+ def tp_size(self) -> int:
+ return self.remote_tp_size[self.engine_id]
+
+ @property
+ def block_size(self) -> int:
+ return self.remote_block_size[self.engine_id]
def tp_ratio(
self,
@@ -957,13 +963,12 @@ class NixlConnectorWorker:
self.xfer_stats = NixlKVConnectorStats()
self.kv_topo = self.TpKVTopology(
- tp_size=self.world_size,
tp_rank=self.tp_rank,
+ engine_id=self.engine_id,
remote_tp_size=self._tp_size, # shared state
+ remote_block_size=self._block_size, # shared state
is_mla=self.use_mla,
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
- block_size=self.block_size,
- remote_block_size=self._block_size,
attn_backend=backend,
)
self._use_pallas = self.kv_topo._use_pallas
@@ -1037,10 +1042,12 @@ class NixlConnectorWorker:
NOT directly supported by NIXL (e.g., tpu)
"""
xfer_buffers: dict[str, torch.Tensor] = {}
+ inv_order = [0, 1, 3, 2, 4]
try:
for layer_name, kv_cache in kv_caches.items():
kv_shape = kv_cache.shape
kv_dtype = kv_cache.dtype
+ permute_shape = False
if (
self.kv_cache_layout == "NHD"
and self.vllm_config.kv_transfer_config is not None
@@ -1054,10 +1061,20 @@ class NixlConnectorWorker:
# Since NHD will not support Decode/Prefill TP_ratio > 1,
# we can leverage host_buffer for permute
self.host_buffer_kv_cache_layout = "HND"
- kv_shape = tuple(kv_shape[i] for i in [0, 1, 3, 2, 4])
+ kv_shape = (
+ tuple(kv_shape[i] for i in inv_order)
+ if not self.use_mla
+ else kv_shape
+ )
+ permute_shape = not self.use_mla
+
xfer_buffers[layer_name] = torch.empty(
kv_shape, dtype=kv_dtype, device="cpu"
)
+ if permute_shape:
+ xfer_buffers[layer_name] = xfer_buffers[layer_name].permute(
+ inv_order
+ )
except MemoryError as e:
logger.error("NIXLConnectorWorker gets %s.", e)
raise
@@ -1156,6 +1173,14 @@ class NixlConnectorWorker:
# to better exploit the memory layout (ie num_blocks is the first dim).
split_k_and_v = self.kv_topo.split_k_and_v
tensor_size_bytes = None
+
+ # TODO (NickLucche): Get kernel_block_size in a cleaner way
+ # NHD default "view" for non-MLA cache
+ if self.device_type == "cpu":
+ block_size_position = -2
+ else:
+ block_size_position = -2 if self.use_mla else -3
+
# Enable different block lengths for different layers when MLA is used.
self.block_len_per_layer = list[int]()
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
@@ -1170,9 +1195,7 @@ class NixlConnectorWorker:
if base_addr in seen_base_addresses:
continue
- # TODO (NickLucche): Get kernel_block_size in a cleaner way
- # NHD default "view" for non-MLA cache
- kernel_block_size = cache.shape[-2] if self.use_mla else cache.shape[-3]
+ kernel_block_size = cache.shape[block_size_position]
if self.block_size != kernel_block_size:
logger.info_once(
@@ -1185,6 +1208,7 @@ class NixlConnectorWorker:
self.block_size // kernel_block_size
)
self.block_size = kernel_block_size
+ self._block_size[self.engine_id] = kernel_block_size
seen_base_addresses.append(base_addr)
curr_tensor_size_bytes = cache.numel() * cache.element_size()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
index 582e42cc466ae..8cd09014cab11 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
@@ -4,12 +4,12 @@ from collections import defaultdict
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from itertools import islice
-from typing import Any
+from typing import Any, ClassVar
import torch
-from vllm.attention import AttentionMetadata
-from vllm.config import VllmConfig
+from vllm.attention import Attention, AttentionBackend, AttentionMetadata
+from vllm.config import VllmConfig, get_layers_from_vllm_config
from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
from vllm.distributed.kv_transfer.kv_connector.v1 import (
KVConnectorBase_V1,
@@ -42,6 +42,8 @@ class OffloadingConnectorMetadata(KVConnectorMetadata):
class OffloadingConnector(KVConnectorBase_V1):
+ prefer_cross_layer_blocks: ClassVar[bool] = True
+
def __init__(
self,
vllm_config: VllmConfig,
@@ -63,6 +65,12 @@ class OffloadingConnector(KVConnectorBase_V1):
assert self.connector_worker is not None
self.connector_worker.register_kv_caches(kv_caches)
+ def register_cross_layers_kv_cache(
+ self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
+ ):
+ assert self.connector_worker is not None
+ self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)
+
def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
assert self.connector_worker is not None
assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
@@ -422,10 +430,35 @@ class OffloadingConnectorWorker:
self._job_counter = job_id + 1
return job_id
- def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
- for src_cls, dst_cls, handler in self.spec.get_handlers(kv_caches):
+ def _register_handlers(
+ self,
+ kv_caches: dict[str, torch.Tensor],
+ attn_backends: dict[str, type[AttentionBackend]],
+ ):
+ for src_cls, dst_cls, handler in self.spec.get_handlers(
+ kv_caches, attn_backends
+ ):
self.worker.register_handler(src_cls, dst_cls, handler)
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+ layer_names = list(kv_caches.keys())
+ layers = get_layers_from_vllm_config(
+ self.spec.vllm_config, Attention, layer_names
+ )
+ attn_backends = {
+ layer_name: layers[layer_name].get_attn_backend()
+ for layer_name in layer_names
+ }
+ self._register_handlers(kv_caches, attn_backends)
+
+ def register_cross_layers_kv_cache(
+ self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
+ ):
+ cross_layer_name = "ALL_LAYERS"
+ kv_caches = {cross_layer_name: kv_cache}
+ attn_backends = {cross_layer_name: attn_backend}
+ self._register_handlers(kv_caches, attn_backends)
+
def start_load_kv(self, metadata: OffloadingConnectorMetadata):
for req_id, transfer_spec in metadata.reqs_to_load.items():
job_id = self._generate_job_id()
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index 852c4c644433f..f81612fd1f4a3 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -1098,6 +1098,12 @@ get_context_model_parallel_group = get_dcp_group
_PP: GroupCoordinator | None = None
+
+def get_pp_group() -> GroupCoordinator:
+ assert _PP is not None, "pipeline model parallel group is not initialized"
+ return _PP
+
+
_DP: GroupCoordinator | None = None
@@ -1114,9 +1120,12 @@ def get_ep_group() -> GroupCoordinator:
return _EP
-def get_pp_group() -> GroupCoordinator:
- assert _PP is not None, "pipeline model parallel group is not initialized"
- return _PP
+_PCP: GroupCoordinator | None = None
+
+
+def get_pcp_group() -> GroupCoordinator:
+ assert _PCP is not None, "prefill context parallel group is not initialized"
+ return _PCP
@deprecated(
@@ -1276,6 +1285,7 @@ def init_distributed_environment(
def initialize_model_parallel(
tensor_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
+ prefill_context_model_parallel_size: int = 1,
decode_context_model_parallel_size: int | None = 1,
backend: str | None = None,
) -> None:
@@ -1325,7 +1335,11 @@ def initialize_model_parallel(
# to get group_ranks for each dimension, transpose that dimension to the
# last dimension, then reshape to 2D, then unbind the last dimension
all_ranks = torch.arange(world_size).reshape(
- -1, data_parallel_size, pipeline_model_parallel_size, tensor_model_parallel_size
+ -1,
+ data_parallel_size,
+ pipeline_model_parallel_size,
+ prefill_context_model_parallel_size,
+ tensor_model_parallel_size,
) # noqa
# Build the tensor model-parallel groups.
@@ -1360,11 +1374,23 @@ def initialize_model_parallel(
group_name="dcp",
)
+ global _PCP
+ assert _PCP is None, "prefill context parallel group is already initialized"
+ group_ranks = (
+ all_ranks.transpose(3, 4)
+ .reshape(-1, prefill_context_model_parallel_size)
+ .unbind(0)
+ )
+ group_ranks = [x.tolist() for x in group_ranks]
+ _PCP = init_model_parallel_group(
+ group_ranks, get_world_group().local_rank, backend, group_name="pcp"
+ )
+
# Build the pipeline model-parallel groups.
global _PP
assert _PP is None, "pipeline model parallel group is already initialized"
group_ranks = (
- all_ranks.transpose(2, 3).reshape(-1, pipeline_model_parallel_size).unbind(0)
+ all_ranks.transpose(2, 4).reshape(-1, pipeline_model_parallel_size).unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
_PP = init_model_parallel_group(
@@ -1373,7 +1399,7 @@ def initialize_model_parallel(
global _DP
assert _DP is None, "data parallel group is already initialized"
- group_ranks = all_ranks.transpose(1, 3).reshape(-1, data_parallel_size).unbind(0)
+ group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
_DP = init_model_parallel_group(
group_ranks, get_world_group().local_rank, backend, group_name="dp"
@@ -1383,7 +1409,12 @@ def initialize_model_parallel(
assert _EP is None, "expert parallel group is already initialized"
group_ranks = (
all_ranks.transpose(1, 2)
- .reshape(-1, data_parallel_size * tensor_model_parallel_size)
+ .reshape(
+ -1,
+ data_parallel_size
+ * prefill_context_model_parallel_size
+ * tensor_model_parallel_size,
+ )
.unbind(0)
)
group_ranks = [x.tolist() for x in group_ranks]
@@ -1393,11 +1424,13 @@ def initialize_model_parallel(
logger.info_once(
"rank %s in world size %s is assigned as "
- "DP rank %s, PP rank %s, TP rank %s, EP rank %s",
+ "DP rank %s, PP rank %s, PCP rank %s, "
+ "TP rank %s, EP rank %s",
rank,
world_size,
_DP.rank_in_group,
_PP.rank_in_group,
+ _PCP.rank_in_group,
_TP.rank_in_group,
_EP.rank_in_group,
)
@@ -1406,6 +1439,7 @@ def initialize_model_parallel(
def ensure_model_parallel_initialized(
tensor_model_parallel_size: int,
pipeline_model_parallel_size: int,
+ prefill_context_model_parallel_size: int = 1,
decode_context_model_parallel_size: int | None = 1,
backend: str | None = None,
) -> None:
@@ -1418,6 +1452,7 @@ def ensure_model_parallel_initialized(
initialize_model_parallel(
tensor_model_parallel_size,
pipeline_model_parallel_size,
+ prefill_context_model_parallel_size,
decode_context_model_parallel_size,
backend,
)
@@ -1434,6 +1469,12 @@ def ensure_model_parallel_initialized(
f"got: {pp_world_size=} vs. "
f"wanted: {pipeline_model_parallel_size=}"
)
+ pcp_world_size = get_pcp_group().world_size
+ assert pcp_world_size == prefill_context_model_parallel_size, (
+ "prefill context parallel group already initialized, but of unexpected size: "
+ f"{pcp_world_size=} vs. "
+ f"{prefill_context_model_parallel_size=}"
+ )
def prepare_communication_buffer_for_model(model: torch.nn.Module):
@@ -1445,6 +1486,8 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module):
"""
if _TP is not None:
_TP.prepare_communication_buffer_for_model(model)
+ if _PCP is not None:
+ _PCP.prepare_communication_buffer_for_model(model)
if _PP is not None:
_PP.prepare_communication_buffer_for_model(model)
if _DP is not None:
@@ -1520,16 +1563,21 @@ def destroy_model_parallel():
_TP.destroy()
_TP = None
- global _PP
- if _PP:
- _PP.destroy()
- _PP = None
-
global _DCP
if _DCP:
_DCP.destroy()
_DCP = None
+ global _PCP
+ if _PCP:
+ _PCP.destroy()
+ _PCP = None
+
+ global _PP
+ if _PP:
+ _PP.destroy()
+ _PP = None
+
global _DP
if _DP:
_DP.destroy()
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
index d011dfdbfbb2e..b7c8f56e18c52 100644
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -367,7 +367,7 @@ class EngineArgs:
config_format: str = ModelConfig.config_format
dtype: ModelDType = ModelConfig.dtype
kv_cache_dtype: CacheDType = CacheConfig.cache_dtype
- seed: int | None = ModelConfig.seed
+ seed: int | None = 0
max_model_len: int | None = ModelConfig.max_model_len
cuda_graph_sizes: list[int] | None = CompilationConfig.cudagraph_capture_sizes
cudagraph_capture_sizes: list[int] | None = (
@@ -389,8 +389,10 @@ class EngineArgs:
nnodes: int = ParallelConfig.nnodes
node_rank: int = ParallelConfig.node_rank
tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
+ prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
+ cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
data_parallel_size: int = ParallelConfig.data_parallel_size
data_parallel_rank: int | None = None
data_parallel_start_rank: int | None = None
@@ -423,7 +425,7 @@ class EngineArgs:
ParallelConfig.max_parallel_loading_workers
)
block_size: BlockSize | None = CacheConfig.block_size
- enable_prefix_caching: bool | None = CacheConfig.enable_prefix_caching
+ enable_prefix_caching: bool | None = None
prefix_caching_hash_algo: PrefixCachingHashAlgo = (
CacheConfig.prefix_caching_hash_algo
)
@@ -482,11 +484,9 @@ class EngineArgs:
fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras
max_cpu_loras: int | None = LoRAConfig.max_cpu_loras
lora_dtype: str | torch.dtype | None = LoRAConfig.lora_dtype
- lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size
ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight
num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override
- num_lookahead_slots: int = SchedulerConfig.num_lookahead_slots
model_loader_extra_config: dict = get_field(LoadConfig, "model_loader_extra_config")
ignore_patterns: str | list[str] = get_field(LoadConfig, "ignore_patterns")
@@ -770,6 +770,15 @@ class EngineArgs:
"--dcp-kv-cache-interleave-size",
**parallel_kwargs["dcp_kv_cache_interleave_size"],
)
+ parallel_group.add_argument(
+ "--cp-kv-cache-interleave-size",
+ **parallel_kwargs["cp_kv_cache_interleave_size"],
+ )
+ parallel_group.add_argument(
+ "--prefill-context-parallel-size",
+ "-pcp",
+ **parallel_kwargs["prefill_context_parallel_size"],
+ )
parallel_group.add_argument(
"--data-parallel-size", "-dp", **parallel_kwargs["data_parallel_size"]
)
@@ -1001,9 +1010,6 @@ class EngineArgs:
)
lora_group.add_argument("--max-loras", **lora_kwargs["max_loras"])
lora_group.add_argument("--max-lora-rank", **lora_kwargs["max_lora_rank"])
- lora_group.add_argument(
- "--lora-extra-vocab-size", **lora_kwargs["lora_extra_vocab_size"]
- )
lora_group.add_argument(
"--lora-dtype",
**lora_kwargs["lora_dtype"],
@@ -1046,10 +1052,18 @@ class EngineArgs:
description=SchedulerConfig.__doc__,
)
scheduler_group.add_argument(
- "--max-num-batched-tokens", **scheduler_kwargs["max_num_batched_tokens"]
+ "--max-num-batched-tokens",
+ **{
+ **scheduler_kwargs["max_num_batched_tokens"],
+ "default": None,
+ },
)
scheduler_group.add_argument(
- "--max-num-seqs", **scheduler_kwargs["max_num_seqs"]
+ "--max-num-seqs",
+ **{
+ **scheduler_kwargs["max_num_seqs"],
+ "default": None,
+ },
)
scheduler_group.add_argument(
"--max-num-partial-prefills", **scheduler_kwargs["max_num_partial_prefills"]
@@ -1062,16 +1076,17 @@ class EngineArgs:
"--long-prefill-token-threshold",
**scheduler_kwargs["long_prefill_token_threshold"],
)
- scheduler_group.add_argument(
- "--num-lookahead-slots", **scheduler_kwargs["num_lookahead_slots"]
- )
# multi-step scheduling has been removed; corresponding arguments
# are no longer supported.
scheduler_group.add_argument(
"--scheduling-policy", **scheduler_kwargs["policy"]
)
scheduler_group.add_argument(
- "--enable-chunked-prefill", **scheduler_kwargs["enable_chunked_prefill"]
+ "--enable-chunked-prefill",
+ **{
+ **scheduler_kwargs["enable_chunked_prefill"],
+ "default": None,
+ },
)
scheduler_group.add_argument(
"--disable-chunked-mm-input", **scheduler_kwargs["disable_chunked_mm_input"]
@@ -1173,29 +1188,52 @@ class EngineArgs:
if check_gguf_file(self.model):
self.quantization = self.load_format = "gguf"
+ # NOTE(woosuk): In V1, we use separate processes for workers (unless
+ # VLLM_ENABLE_V1_MULTIPROCESSING=0), so setting a seed here
+ # doesn't affect the user process.
+ if self.seed is None:
+ logger.warning_once(
+ "`seed=None` is equivalent to `seed=0` in V1 Engine. "
+ "You will no longer be allowed to pass `None` in v0.13.",
+ scope="local",
+ )
+
+ self.seed = 0
+ if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
+ logger.warning(
+ "The global random seed is set to %d. Since "
+ "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
+ "affect the random state of the Python process that "
+ "launched vLLM.",
+ self.seed,
+ )
+
if self.disable_mm_preprocessor_cache:
- logger.warning(
+ logger.warning_once(
"`--disable-mm-preprocessor-cache` is deprecated "
"and will be removed in v0.13. "
"Please use `--mm-processor-cache-gb 0` instead.",
+ scope="local",
)
self.mm_processor_cache_gb = 0
elif envs.VLLM_MM_INPUT_CACHE_GIB != 4:
- logger.warning(
+ logger.warning_once(
"VLLM_MM_INPUT_CACHE_GIB` is deprecated "
"and will be removed in v0.13. "
"Please use `--mm-processor-cache-gb %d` instead.",
envs.VLLM_MM_INPUT_CACHE_GIB,
+ scope="local",
)
self.mm_processor_cache_gb = envs.VLLM_MM_INPUT_CACHE_GIB
if self.enable_multimodal_encoder_data_parallel:
- logger.warning(
+ logger.warning_once(
"--enable-multimodal-encoder-data-parallel` is deprecated "
"and will be removed in v0.13. "
- "Please use `--mm-encoder-tp-mode data` instead."
+ "Please use `--mm-encoder-tp-mode data` instead.",
+ scope="local",
)
self.mm_encoder_tp_mode = "data"
@@ -1354,11 +1392,10 @@ class EngineArgs:
# Set default arguments for V1 Engine.
self._set_default_args(usage_context, model_config)
# Disable chunked prefill and prefix caching for:
- # POWER (ppc64le)/ARM/s390x/RISCV CPUs in V1
+ # POWER (ppc64le)/s390x/RISCV CPUs in V1
if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
CpuArchEnum.POWERPC,
CpuArchEnum.S390X,
- CpuArchEnum.ARM,
CpuArchEnum.RISCV,
):
logger.info(
@@ -1488,7 +1525,7 @@ class EngineArgs:
# Local DP rank = 1, use pure-external LB.
if data_parallel_external_lb:
assert self.data_parallel_rank is not None, (
- "data_parallel_rank or node_rank must be spefified if "
+ "data_parallel_rank or node_rank must be specified if "
"data_parallel_external_lb is enable."
)
assert self.data_parallel_size_local in (1, None), (
@@ -1588,6 +1625,7 @@ class EngineArgs:
parallel_config = ParallelConfig(
pipeline_parallel_size=self.pipeline_parallel_size,
tensor_parallel_size=self.tensor_parallel_size,
+ prefill_context_parallel_size=self.prefill_context_parallel_size,
data_parallel_size=self.data_parallel_size,
data_parallel_rank=self.data_parallel_rank or 0,
data_parallel_external_lb=data_parallel_external_lb,
@@ -1619,6 +1657,7 @@ class EngineArgs:
worker_extension_cls=self.worker_extension_cls,
decode_context_parallel_size=self.decode_context_parallel_size,
dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
_api_process_count=self._api_process_count,
_api_process_rank=self._api_process_rank,
)
@@ -1628,18 +1667,11 @@ class EngineArgs:
target_parallel_config=parallel_config,
)
- # make sure num_lookahead_slots is set appropriately depending on
- # whether speculative decoding is enabled
- num_lookahead_slots = self.num_lookahead_slots
- if speculative_config is not None:
- num_lookahead_slots = speculative_config.num_lookahead_slots
-
scheduler_config = SchedulerConfig(
runner_type=model_config.runner_type,
max_num_batched_tokens=self.max_num_batched_tokens,
max_num_seqs=self.max_num_seqs,
max_model_len=model_config.max_model_len,
- num_lookahead_slots=num_lookahead_slots,
enable_chunked_prefill=self.enable_chunked_prefill,
disable_chunked_mm_input=self.disable_chunked_mm_input,
is_multimodal_model=model_config.is_multimodal_model,
@@ -1666,7 +1698,6 @@ class EngineArgs:
max_loras=self.max_loras,
default_mm_loras=self.default_mm_loras,
fully_sharded_loras=self.fully_sharded_loras,
- lora_extra_vocab_size=self.lora_extra_vocab_size,
lora_dtype=self.lora_dtype,
max_cpu_loras=self.max_cpu_loras
if self.max_cpu_loras and self.max_cpu_loras > 0
@@ -1940,6 +1971,16 @@ class EngineArgs:
default_prefix_caching,
) = self.get_chunked_prefill_prefix_caching_defaults(model_config)
+ if self.prefill_context_parallel_size > 1:
+ default_chunked_prefill = False
+ default_prefix_caching = False
+ logger.warning_once(
+ "--prefill-context-parallel-size > 1 is not compatible with "
+ "chunked prefill and prefix caching now. Chunked prefill "
+ "and prefix caching have been disabled by default.",
+ scope="local",
+ )
+
if self.enable_chunked_prefill is None:
self.enable_chunked_prefill = default_chunked_prefill
@@ -1947,15 +1988,27 @@ class EngineArgs:
"%s chunked prefill by default",
"Enabling" if default_chunked_prefill else "Disabling",
)
+ elif (
+ model_config.runner_type == "generate"
+ and not self.enable_chunked_prefill
+ and default_chunked_prefill
+ ):
+ logger.warning_once(
+ "This model does not officially support disabling chunked prefill. "
+ "Disabling this manually may cause the engine to crash "
+ "or produce incorrect outputs.",
+ scope="local",
+ )
elif (
model_config.runner_type == "pooling"
and self.enable_chunked_prefill
and not default_chunked_prefill
):
- logger.warning(
+ logger.warning_once(
"This model does not officially support chunked prefill. "
"Enabling this manually may cause the engine to crash "
"or produce incorrect outputs.",
+ scope="local",
)
if self.enable_prefix_caching is None:
@@ -1970,10 +2023,11 @@ class EngineArgs:
and self.enable_prefix_caching
and not default_prefix_caching
):
- logger.warning(
+ logger.warning_once(
"This model does not officially support prefix caching. "
"Enabling this manually may cause the engine to crash "
"or produce incorrect outputs.",
+ scope="local",
)
world_size = self.pipeline_parallel_size * self.tensor_parallel_size
diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py
index 462d2c4e50e73..5e3374f9f6a10 100644
--- a/vllm/engine/protocol.py
+++ b/vllm/engine/protocol.py
@@ -149,6 +149,33 @@ class EngineClient(ABC):
"""Load a new LoRA adapter into the engine for future requests."""
...
+ @abstractmethod
+ async def pause_generation(
+ self,
+ *,
+ wait_for_inflight_requests: bool = False,
+ clear_cache: bool = True,
+ ) -> None:
+ """Pause new generation/encoding requests.
+
+ Args:
+ wait_for_inflight_requests: When ``True`` waits for in-flight requests
+ to finish before pausing. When ``False`` (default), aborts in-flight
+ requests immediately.
+ clear_cache: Whether to clear KV and prefix caches after draining.
+ """
+ ...
+
+ @abstractmethod
+ async def resume_generation(self) -> None:
+ """Resume accepting generation/encoding requests."""
+ ...
+
+ @abstractmethod
+ async def is_paused(self) -> bool:
+ """Return whether the engine is currently paused."""
+ ...
+
async def scale_elastic_ep(
self, new_data_parallel_size: int, drain_timeout: int = 300
) -> None:
diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py
index 3b722c2d92770..bf80856c1bbfc 100644
--- a/vllm/entrypoints/chat_utils.py
+++ b/vllm/entrypoints/chat_utils.py
@@ -94,6 +94,22 @@ class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):
"""
+class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False):
+ audio_embeds: str | dict[str, str] | None
+ """
+ The audio embeddings. It can be either:
+ - A single base64 string representing a serialized torch tensor.
+ - A dictionary where each value is a base64 string.
+ """
+ type: Required[Literal["audio_embeds"]]
+ """The type of the content part."""
+ uuid: str | None
+ """
+ User-provided UUID of a media. User must guarantee that it is properly
+ generated and unique for different medias.
+ """
+
+
class VideoURL(TypedDict, total=False):
url: Required[str]
"""
@@ -211,6 +227,7 @@ ChatCompletionContentPartParam: TypeAlias = (
| CustomChatCompletionContentPILImageParam
| CustomChatCompletionContentSimpleImageParam
| ChatCompletionContentPartImageEmbedsParam
+ | ChatCompletionContentPartAudioEmbedsParam
| CustomChatCompletionContentSimpleAudioParam
| CustomChatCompletionContentSimpleVideoParam
| str
@@ -599,7 +616,7 @@ def resolve_chat_template_content_format(
return detected_format
-ModalityStr = Literal["image", "audio", "video", "image_embeds"]
+ModalityStr = Literal["image", "audio", "video", "image_embeds", "audio_embeds"]
_T = TypeVar("_T")
@@ -684,6 +701,11 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
mm_uuids["image"] = uuids_by_modality["image_embeds"]
if "image" in uuids_by_modality:
mm_uuids["image"] = uuids_by_modality["image"] # UUIDs of images
+ if "audio_embeds" in uuids_by_modality:
+ audio_embeds_uuids = uuids_by_modality["audio_embeds"]
+ if len(audio_embeds_uuids) > 1:
+ raise ValueError("Only one message can have {'type': 'audio_embeds'}")
+ mm_uuids["audio"] = uuids_by_modality["audio_embeds"]
if "audio" in uuids_by_modality:
mm_uuids["audio"] = uuids_by_modality["audio"] # UUIDs of audios
if "video" in uuids_by_modality:
@@ -703,6 +725,8 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
items_by_modality = dict(self._items_by_modality)
if "image" in items_by_modality and "image_embeds" in items_by_modality:
raise ValueError("Mixing raw image and embedding inputs is not allowed")
+ if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
+ raise ValueError("Mixing raw audio and embedding inputs is not allowed")
if "image_embeds" in items_by_modality:
image_embeds_lst = items_by_modality["image_embeds"]
@@ -711,6 +735,11 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
mm_inputs["image"] = image_embeds_lst[0]
if "image" in items_by_modality:
mm_inputs["image"] = items_by_modality["image"] # A list of images
+ if "audio_embeds" in items_by_modality:
+ audio_embeds_lst = items_by_modality["audio_embeds"]
+ if len(audio_embeds_lst) > 1:
+ raise ValueError("Only one message can have {'type': 'audio_embeds'}")
+ mm_inputs["audio"] = audio_embeds_lst[0]
if "audio" in items_by_modality:
mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
if "video" in items_by_modality:
@@ -738,6 +767,8 @@ class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
if "image" in items_by_modality and "image_embeds" in items_by_modality:
raise ValueError("Mixing raw image and embedding inputs is not allowed")
+ if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
+ raise ValueError("Mixing raw audio and embedding inputs is not allowed")
if "image_embeds" in items_by_modality:
image_embeds_lst = items_by_modality["image_embeds"]
@@ -746,6 +777,11 @@ class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
mm_inputs["image"] = image_embeds_lst[0]
if "image" in items_by_modality:
mm_inputs["image"] = items_by_modality["image"] # A list of images
+ if "audio_embeds" in items_by_modality:
+ audio_embeds_lst = items_by_modality["audio_embeds"]
+ if len(audio_embeds_lst) > 1:
+ raise ValueError("Only one message can have {'type': 'audio_embeds'}")
+ mm_inputs["audio"] = audio_embeds_lst[0]
if "audio" in items_by_modality:
mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
if "video" in items_by_modality:
@@ -804,6 +840,14 @@ class BaseMultiModalContentParser(ABC):
) -> None:
raise NotImplementedError
+ @abstractmethod
+ def parse_audio_embeds(
+ self,
+ audio_embeds: str | dict[str, str] | None,
+ uuid: str | None = None,
+ ) -> None:
+ raise NotImplementedError
+
@abstractmethod
def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
raise NotImplementedError
@@ -861,6 +905,31 @@ class MultiModalContentParser(BaseMultiModalContentParser):
self._add_placeholder("image", placeholder)
+ def parse_audio_embeds(
+ self,
+ audio_embeds: str | dict[str, str] | None,
+ uuid: str | None = None,
+ ) -> None:
+ mm_config = self.model_config.get_multimodal_config()
+ if not mm_config.enable_mm_embeds:
+ raise ValueError(
+ "You must set `--enable-mm-embeds` to input `audio_embeds`"
+ )
+
+ if isinstance(audio_embeds, dict):
+ embeds = {
+ k: self._connector.fetch_audio_embedding(v)
+ for k, v in audio_embeds.items()
+ }
+ placeholder = self._tracker.add("audio_embeds", embeds, uuid)
+ elif isinstance(audio_embeds, str):
+ embedding = self._connector.fetch_audio_embedding(audio_embeds)
+ placeholder = self._tracker.add("audio_embeds", embedding, uuid)
+ else:
+ placeholder = self._tracker.add("audio_embeds", None, uuid)
+
+ self._add_placeholder("audio", placeholder)
+
def parse_image_pil(
self, image_pil: Image.Image | None, uuid: str | None = None
) -> None:
@@ -950,6 +1019,67 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):
placeholder = self._tracker.add("image_embeds", future, uuid)
self._add_placeholder("image", placeholder)
+ def parse_audio_embeds(
+ self,
+ audio_embeds: str | dict[str, str] | None,
+ uuid: str | None = None,
+ ) -> None:
+ mm_config = self.model_config.get_multimodal_config()
+ if not mm_config.enable_mm_embeds:
+ raise ValueError(
+ "You must set `--enable-mm-embeds` to input `audio_embeds`"
+ )
+
+ logger.info(
+ "🎵 Parsing audio_embeds: type=%s, uuid=%s, is_dict=%s, "
+ "is_str=%s, is_none=%s",
+ type(audio_embeds).__name__,
+ uuid,
+ isinstance(audio_embeds, dict),
+ isinstance(audio_embeds, str),
+ audio_embeds is None,
+ )
+
+ future: asyncio.Future[str | dict[str, str] | None] = asyncio.Future()
+
+ if isinstance(audio_embeds, dict):
+ logger.info(
+ "🎵 Processing dict audio_embeds with %d entries",
+ len(audio_embeds),
+ )
+ embeds = {
+ k: self._connector.fetch_audio_embedding(v)
+ for k, v in audio_embeds.items()
+ }
+ future.set_result(embeds)
+ logger.info(
+ "🎵 Successfully loaded %d audio embeddings from dict",
+ len(embeds),
+ )
+
+ if isinstance(audio_embeds, str):
+ base64_size = len(audio_embeds)
+ logger.info(
+ "🎵 Processing base64 audio_embeds: %d chars (%.2f KB)",
+ base64_size,
+ base64_size / 1024,
+ )
+ embedding = self._connector.fetch_audio_embedding(audio_embeds)
+ future.set_result(embedding)
+ logger.info(
+ "🎵 Successfully loaded audio embedding tensor: shape=%s, dtype=%s",
+ embedding.shape,
+ embedding.dtype,
+ )
+
+ if audio_embeds is None:
+ logger.info("🎵 Audio embeds is None (UUID-only reference)")
+ future.set_result(None)
+
+ placeholder = self._tracker.add("audio_embeds", future, uuid)
+ self._add_placeholder("audio", placeholder)
+ logger.info("🎵 Added audio_embeds placeholder with uuid=%s", uuid)
+
def parse_image_pil(
self, image_pil: Image.Image | None, uuid: str | None = None
) -> None:
@@ -1132,6 +1262,7 @@ def _get_full_multimodal_text_prompt(
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
+_AudioEmbedsParser = partial(cast, ChatCompletionContentPartAudioEmbedsParam)
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
@@ -1152,9 +1283,11 @@ MM_PARSER_MAP: dict[
"text": lambda part: _TextParser(part).get("text", None),
"thinking": lambda part: _ThinkParser(part).get("thinking", None),
"input_text": lambda part: _TextParser(part).get("text", None),
+ "output_text": lambda part: _TextParser(part).get("text", None),
"input_image": lambda part: _ResponsesInputImageParser(part).get("image_url", None),
"image_url": lambda part: _ImageParser(part).get("image_url", {}).get("url", None),
"image_embeds": lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
+ "audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),
"image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
"audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
"input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),
@@ -1223,8 +1356,17 @@ def _parse_chat_message_content_mm_part(
)
image_embeds = image_params.get("image_embeds", None)
return "image_embeds", image_embeds
+ if "audio_embeds" in part:
+ # "audio_embeds" could be None if UUID is provided.
+ audio_params = cast( # type: ignore[assignment]
+ ChatCompletionContentPartAudioEmbedsParam, part
+ )
+ audio_embeds = audio_params.get("audio_embeds", None)
+ return "audio_embeds", audio_embeds
if "audio_url" in part:
- audio_params = cast(CustomChatCompletionContentSimpleAudioParam, part)
+ audio_params = cast( # type: ignore[assignment]
+ CustomChatCompletionContentSimpleAudioParam, part
+ )
audio_url = audio_params.get("audio_url", None)
if isinstance(audio_url, dict):
# Can potentially happen if user provides a uuid
@@ -1322,7 +1464,7 @@ def _parse_chat_message_content_part(
)
return None
- if part_type in ("text", "input_text", "refusal", "thinking"):
+ if part_type in ("text", "input_text", "output_text", "refusal", "thinking"):
str_content = cast(str, content)
if wrap_dicts:
return {"type": "text", "text": str_content}
@@ -1348,6 +1490,10 @@ def _parse_chat_message_content_part(
content = cast(str | dict[str, str], content) if content is not None else None
mm_parser.parse_image_embeds(content, uuid)
modality = "image"
+ elif part_type == "audio_embeds":
+ content = cast(str | dict[str, str], content) if content is not None else None
+ mm_parser.parse_audio_embeds(content, uuid)
+ modality = "audio"
elif part_type == "audio_url":
str_content = cast(str, content)
mm_parser.parse_audio(str_content, uuid)
@@ -1437,7 +1583,8 @@ def _postprocess_messages(messages: list[ConversationMessage]) -> None:
for item in message["tool_calls"]:
# if arguments is None or empty string, set to {}
if content := item["function"].get("arguments"):
- item["function"]["arguments"] = json.loads(content)
+ if not isinstance(content, (dict, list)):
+ item["function"]["arguments"] = json.loads(content)
else:
item["function"]["arguments"] = {}
diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
index b0786bd355aa6..848916dbd8763 100644
--- a/vllm/entrypoints/llm.py
+++ b/vllm/entrypoints/llm.py
@@ -339,7 +339,6 @@ class LLM:
log_non_default_args(engine_args)
- # Create the Engine (autoselects V0 vs V1)
self.llm_engine = LLMEngine.from_engine_args(
engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
)
@@ -466,7 +465,7 @@ class LLM:
):
return lora_request
- if not isinstance(prompts, Sequence):
+ if not isinstance(prompts, Sequence) or isinstance(prompts, str):
prompts = [prompts]
optional_loras = (
diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py
index 3cf66fcd27e2a..70174250ceabe 100644
--- a/vllm/entrypoints/openai/api_server.py
+++ b/vllm/entrypoints/openai/api_server.py
@@ -5,6 +5,7 @@ import hashlib
import importlib
import inspect
import json
+import logging
import multiprocessing
import multiprocessing.forkserver as forkserver
import os
@@ -393,6 +394,84 @@ async def get_server_load_metrics(request: Request):
return JSONResponse(content={"server_load": request.app.state.server_load_metrics})
+@router.post("/pause")
+async def pause_generation(
+ raw_request: Request,
+ wait_for_inflight_requests: bool = Query(False),
+ clear_cache: bool = Query(True),
+) -> JSONResponse:
+ """Pause generation requests to allow weight updates.
+
+ Args:
+ wait_for_inflight_requests: When ``True`` waits for in-flight
+ requests to finish before pausing. When ``False`` (default),
+ aborts any in-flight requests immediately.
+ clear_cache: Whether to clear KV/prefix caches after draining.
+ """
+
+ engine = engine_client(raw_request)
+
+ try:
+ await engine.pause_generation(
+ wait_for_inflight_requests=wait_for_inflight_requests,
+ clear_cache=clear_cache,
+ )
+ return JSONResponse(
+ content={"status": "paused"},
+ status_code=HTTPStatus.OK.value,
+ )
+
+ except ValueError as err:
+ return JSONResponse(
+ content={"error": str(err)},
+ status_code=HTTPStatus.BAD_REQUEST.value,
+ )
+ except Exception as err: # pragma: no cover - defensive
+ logger.exception("Failed to pause generation")
+ return JSONResponse(
+ content={"error": f"Failed to pause generation: {err}"},
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+ )
+
+
+@router.post("/resume")
+async def resume_generation(raw_request: Request) -> JSONResponse:
+ """Resume generation after a pause."""
+
+ engine = engine_client(raw_request)
+
+ try:
+ await engine.resume_generation()
+ return JSONResponse(
+ content={"status": "resumed"},
+ status_code=HTTPStatus.OK.value,
+ )
+ except Exception as err: # pragma: no cover - defensive
+ logger.exception("Failed to resume generation")
+ return JSONResponse(
+ content={"error": f"Failed to resume generation: {err}"},
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+ )
+
+
+@router.get("/is_paused")
+async def is_paused(raw_request: Request) -> JSONResponse:
+ """Return the current pause status."""
+
+ engine = engine_client(raw_request)
+
+ try:
+ paused = await engine.is_paused()
+ except Exception as err: # pragma: no cover - defensive
+ logger.exception("Failed to fetch pause status")
+ return JSONResponse(
+ content={"error": f"Failed to fetch pause status: {err}"},
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+ )
+
+ return JSONResponse(content={"is_paused": paused})
+
+
@router.post(
"/tokenize",
dependencies=[Depends(validate_json_request)],
@@ -2020,6 +2099,9 @@ async def run_server(args, **uvicorn_kwargs) -> None:
# Add process-specific prefix to stdout and stderr.
decorate_logs("APIServer")
+ # Suppress verbose logs from model_hosting_container_standards
+ logging.getLogger("model_hosting_container_standards").setLevel(logging.ERROR)
+
listen_address, sock = setup_server(args)
await run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
diff --git a/vllm/entrypoints/openai/protocol.py b/vllm/entrypoints/openai/protocol.py
index 65bd15ba387b9..b352c3ad01db0 100644
--- a/vllm/entrypoints/openai/protocol.py
+++ b/vllm/entrypoints/openai/protocol.py
@@ -377,7 +377,7 @@ class ResponsesRequest(OpenAIBaseModel):
"environments. The salt should be random, protected from "
"access by 3rd parties, and long enough to be "
"unpredictable (e.g., 43 characters base64-encoded, corresponding "
- "to 256 bit). Not supported by vLLM engine V0."
+ "to 256 bit)."
),
)
@@ -565,7 +565,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
user: str | None = None
# --8<-- [start:chat-completion-sampling-params]
- best_of: int | None = None
use_beam_search: bool = False
top_k: int | None = None
min_p: float | None = None
@@ -764,7 +763,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
"environments. The salt should be random, protected from "
"access by 3rd parties, and long enough to be "
"unpredictable (e.g., 43 characters base64-encoded, corresponding "
- "to 256 bit). Not supported by vLLM engine V0."
+ "to 256 bit)."
),
)
kv_transfer_params: dict[str, Any] | None = Field(
@@ -889,7 +888,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
extra_args["kv_transfer_params"] = self.kv_transfer_params
return SamplingParams.from_optional(
n=self.n,
- best_of=self.best_of,
presence_penalty=self.presence_penalty,
frequency_penalty=self.frequency_penalty,
repetition_penalty=repetition_penalty,
@@ -1088,7 +1086,6 @@ class CompletionRequest(OpenAIBaseModel):
# https://platform.openai.com/docs/api-reference/completions/create
model: str | None = None
prompt: list[int] | list[list[int]] | str | list[str] | None = None
- best_of: int | None = None
echo: bool | None = False
frequency_penalty: float | None = 0.0
logit_bias: dict[str, float] | None = None
@@ -1252,7 +1249,7 @@ class CompletionRequest(OpenAIBaseModel):
"environments. The salt should be random, protected from "
"access by 3rd parties, and long enough to be "
"unpredictable (e.g., 43 characters base64-encoded, corresponding "
- "to 256 bit). Not supported by vLLM engine V0."
+ "to 256 bit)."
),
)
@@ -1375,7 +1372,6 @@ class CompletionRequest(OpenAIBaseModel):
extra_args["kv_transfer_params"] = self.kv_transfer_params
return SamplingParams.from_optional(
n=self.n,
- best_of=self.best_of,
presence_penalty=self.presence_penalty,
frequency_penalty=self.frequency_penalty,
repetition_penalty=repetition_penalty,
diff --git a/vllm/entrypoints/openai/serving_chat.py b/vllm/entrypoints/openai/serving_chat.py
index 59e1c8d531793..6cc685acd6728 100644
--- a/vllm/entrypoints/openai/serving_chat.py
+++ b/vllm/entrypoints/openai/serving_chat.py
@@ -319,6 +319,7 @@ class OpenAIServingChat(OpenAIServing):
request_id=request_id,
params=sampling_params,
lora_request=lora_request,
+ trace_headers=trace_headers,
)
else:
engine_request, tokenization_kwargs = await self._process_inputs(
diff --git a/vllm/entrypoints/openai/serving_completion.py b/vllm/entrypoints/openai/serving_completion.py
index a114b77ebc16b..9681aa8c71e6d 100644
--- a/vllm/entrypoints/openai/serving_completion.py
+++ b/vllm/entrypoints/openai/serving_completion.py
@@ -216,6 +216,7 @@ class OpenAIServingCompletion(OpenAIServing):
request_id=request_id,
params=sampling_params,
lora_request=lora_request,
+ trace_headers=trace_headers,
)
else:
engine_request, tokenization_kwargs = await self._process_inputs(
@@ -249,14 +250,8 @@ class OpenAIServingCompletion(OpenAIServing):
model_name = self.models.model_name(lora_request)
num_prompts = len(engine_prompts)
- # Similar to the OpenAI API, when n != best_of, we do not stream the
- # results. Noting that best_of is only supported in V0. In addition,
- # we do not stream the results when use beam search.
- stream = (
- request.stream
- and (request.best_of is None or request.n == request.best_of)
- and not request.use_beam_search
- )
+ # We do not stream the results when using beam search.
+ stream = request.stream and not request.use_beam_search
# Streaming response
if stream:
diff --git a/vllm/entrypoints/openai/serving_engine.py b/vllm/entrypoints/openai/serving_engine.py
index c50b0c4a23e17..7dab5dbacd28c 100644
--- a/vllm/entrypoints/openai/serving_engine.py
+++ b/vllm/entrypoints/openai/serving_engine.py
@@ -10,6 +10,7 @@ from concurrent.futures import ThreadPoolExecutor
from http import HTTPStatus
from typing import Any, ClassVar, Generic, TypeAlias, TypeVar
+import numpy as np
import torch
from fastapi import Request
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
@@ -342,6 +343,7 @@ class OpenAIServing:
request_id: str,
params: BeamSearchParams,
lora_request: LoRARequest | None = None,
+ trace_headers: Mapping[str, str] | None = None,
) -> AsyncGenerator[RequestOutput, None]:
beam_width = params.beam_width
max_tokens = params.max_tokens
@@ -389,8 +391,9 @@ class OpenAIServing:
sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty)
+ logprobs_num = 2 * beam_width
beam_search_params = SamplingParams(
- logprobs=2 * beam_width,
+ logprobs=logprobs_num,
max_tokens=1,
temperature=temperature,
)
@@ -435,6 +438,7 @@ class OpenAIServing:
beam_search_params,
request_id_item,
lora_request=lora_req,
+ trace_headers=trace_headers,
)
)
)
@@ -443,40 +447,75 @@ class OpenAIServing:
output = [x[0] for x in await asyncio.gather(*tasks)]
new_beams = []
- for i, current_beam in enumerate(all_beams):
- result = output[i]
-
+ # Store all new tokens generated by beam
+ all_beams_token_id = []
+ # Store the cumulative probability of all tokens
+ # generated by beam search
+ all_beams_logprob = []
+ # Iterate through all beam inference results
+ for i, result in enumerate(output):
+ current_beam = all_beams[i]
if result.outputs[0].logprobs is not None:
logprobs = result.outputs[0].logprobs[0]
- for token_id, logprob_obj in logprobs.items():
- if token_id == eos_token_id and not ignore_eos:
- completed.append(
- BeamSearchSequence(
- tokens=current_beam.tokens + [token_id]
- if include_stop_str_in_output
- else current_beam.tokens,
- logprobs=current_beam.logprobs + [logprobs],
- cum_logprob=current_beam.cum_logprob
- + logprob_obj.logprob,
- finish_reason="stop",
- stop_reason=eos_token_id,
- )
- )
- else:
- new_beams.append(
- BeamSearchSequence(
- tokens=current_beam.tokens + [token_id],
- logprobs=current_beam.logprobs + [logprobs],
- lora_request=current_beam.lora_request,
- cum_logprob=current_beam.cum_logprob
- + logprob_obj.logprob,
- multi_modal_data=current_beam.multi_modal_data,
- mm_processor_kwargs=current_beam.mm_processor_kwargs,
- )
- )
+ all_beams_token_id.extend(list(logprobs.keys()))
+ all_beams_logprob.extend(
+ [
+ current_beam.cum_logprob + obj.logprob
+ for obj in logprobs.values()
+ ]
+ )
- sorted_beams = sorted(new_beams, key=sort_beams_key, reverse=True)
- all_beams = sorted_beams[:beam_width]
+ # Handle the token for the end of sentence (EOS)
+ all_beams_token_id = np.array(all_beams_token_id)
+ all_beams_logprob = np.array(all_beams_logprob)
+
+ if not ignore_eos:
+ # Get the index position of eos token in all generated results
+ eos_idx = np.where(all_beams_token_id == eos_token_id)[0]
+ for idx in eos_idx:
+ current_beam = all_beams[idx // logprobs_num]
+ result = output[idx // logprobs_num]
+ assert result.outputs[0].logprobs is not None
+ logprobs_entry = result.outputs[0].logprobs[0]
+ completed.append(
+ BeamSearchSequence(
+ tokens=current_beam.tokens + [eos_token_id]
+ if include_stop_str_in_output
+ else current_beam.tokens,
+ logprobs=current_beam.logprobs + [logprobs_entry],
+ cum_logprob=float(all_beams_logprob[idx]),
+ finish_reason="stop",
+ stop_reason=eos_token_id,
+ )
+ )
+ # After processing, set the log probability of the eos condition
+ # to negative infinity.
+ all_beams_logprob[eos_idx] = -np.inf
+
+ # Processing non-EOS tokens
+ # Get indices of the top beam_width probabilities
+ topn_idx = np.argpartition(np.negative(all_beams_logprob), beam_width)[
+ :beam_width
+ ]
+
+ for idx in topn_idx:
+ current_beam = all_beams[idx // logprobs_num]
+ result = output[idx // logprobs_num]
+ token_id = int(all_beams_token_id[idx])
+ assert result.outputs[0].logprobs is not None
+ logprobs_entry = result.outputs[0].logprobs[0]
+ new_beams.append(
+ BeamSearchSequence(
+ tokens=current_beam.tokens + [token_id],
+ logprobs=current_beam.logprobs + [logprobs_entry],
+ lora_request=current_beam.lora_request,
+ cum_logprob=float(all_beams_logprob[idx]),
+ multi_modal_data=current_beam.multi_modal_data,
+ mm_processor_kwargs=current_beam.mm_processor_kwargs,
+ )
+ )
+
+ all_beams = new_beams
completed.extend(all_beams)
sorted_completed = sorted(completed, key=sort_beams_key, reverse=True)
diff --git a/vllm/entrypoints/openai/tool_parsers/glm4_moe_tool_parser.py b/vllm/entrypoints/openai/tool_parsers/glm4_moe_tool_parser.py
index 120e63b929b16..389e9754b34da 100644
--- a/vllm/entrypoints/openai/tool_parsers/glm4_moe_tool_parser.py
+++ b/vllm/entrypoints/openai/tool_parsers/glm4_moe_tool_parser.py
@@ -78,7 +78,7 @@ class Glm4MoeModelToolParser(ToolParser):
.get("type", None)
)
return arg_type == "string"
- logger.warning("No tool named '%s'.", tool_name)
+ logger.debug("No tool named '%s'.", tool_name)
return False
def _deserialize(value: str) -> Any:
diff --git a/vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py b/vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py
index a84c9e4547168..2b84c60a3b841 100644
--- a/vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py
+++ b/vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py
@@ -60,7 +60,8 @@ class KimiK2ToolParser(ToolParser):
self.tool_call_end_token: str = "<|tool_call_end|>"
self.tool_call_regex = re.compile(
- r"<\|tool_call_begin\|>\s*(?P.+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P.*?)\s*<\|tool_call_end\|>"
+ r"<\|tool_call_begin\|>\s*(?P[^<]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P(?:(?!<\|tool_call_begin\|>).)*?)\s*<\|tool_call_end\|>",
+ re.DOTALL,
)
self.stream_tool_call_portion_regex = re.compile(
diff --git a/vllm/entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py b/vllm/entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py
index 26261c0065ead..9d4c079eba188 100644
--- a/vllm/entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py
+++ b/vllm/entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py
@@ -128,7 +128,7 @@ class Qwen3CoderToolParser(ToolParser):
return params
else:
return {}
- logger.warning("Tool '%s' is not defined in the tools list.", func_name)
+ logger.debug("Tool '%s' is not defined in the tools list.", func_name)
return {}
def _convert_param_value(
@@ -141,7 +141,7 @@ class Qwen3CoderToolParser(ToolParser):
if param_name not in param_config:
if param_config != {}:
- logger.warning(
+ logger.debug(
"Parsed parameter '%s' is not defined in the tool "
"parameters for tool '%s', directly returning the "
"string value.",
@@ -169,7 +169,7 @@ class Qwen3CoderToolParser(ToolParser):
try:
return int(param_value)
except (ValueError, TypeError):
- logger.warning(
+ logger.debug(
"Parsed value '%s' of parameter '%s' is not an "
"integer in tool '%s', degenerating to string.",
param_value,
@@ -186,7 +186,7 @@ class Qwen3CoderToolParser(ToolParser):
else int(float_param_value)
)
except (ValueError, TypeError):
- logger.warning(
+ logger.debug(
"Parsed value '%s' of parameter '%s' is not a float "
"in tool '%s', degenerating to string.",
param_value,
@@ -197,7 +197,7 @@ class Qwen3CoderToolParser(ToolParser):
elif param_type in ["boolean", "bool", "binary"]:
param_value = param_value.lower()
if param_value not in ["true", "false"]:
- logger.warning(
+ logger.debug(
"Parsed value '%s' of parameter '%s' is not a boolean "
"(`true` or `false`) in tool '%s', degenerating to "
"false.",
@@ -216,7 +216,7 @@ class Qwen3CoderToolParser(ToolParser):
param_value = json.loads(param_value)
return param_value
except (json.JSONDecodeError, TypeError, ValueError):
- logger.warning(
+ logger.debug(
"Parsed value '%s' of parameter '%s' cannot be "
"parsed with json.loads in tool '%s', will try "
"other methods to parse it.",
@@ -227,7 +227,7 @@ class Qwen3CoderToolParser(ToolParser):
try:
param_value = ast.literal_eval(param_value) # safer
except (ValueError, SyntaxError, TypeError):
- logger.warning(
+ logger.debug(
"Parsed value '%s' of parameter '%s' cannot be "
"converted via Python `ast.literal_eval()` in tool "
"'%s', degenerating to string.",
diff --git a/vllm/entrypoints/responses_utils.py b/vllm/entrypoints/responses_utils.py
index d966f58804b67..912e8a690573d 100644
--- a/vllm/entrypoints/responses_utils.py
+++ b/vllm/entrypoints/responses_utils.py
@@ -10,6 +10,7 @@ from openai.types.chat.chat_completion_message_tool_call_param import (
Function as FunctionCallTool,
)
from openai.types.responses import ResponseFunctionToolCall
+from openai.types.responses.response_reasoning_item import ResponseReasoningItem
from openai.types.responses.tool import Tool
from vllm import envs
@@ -37,6 +38,18 @@ def construct_chat_message_with_tool_call(
)
],
)
+ elif isinstance(item, ResponseReasoningItem):
+ reasoning_content = ""
+ if item.encrypted_content:
+ raise ValueError("Encrypted content is not supported.")
+ if len(item.summary) == 1:
+ reasoning_content = item.summary[0].text
+ elif item.content and len(item.content) == 1:
+ reasoning_content = item.content[0].text
+ return {
+ "role": "assistant",
+ "reasoning": reasoning_content,
+ }
elif item.get("type") == "function_call_output":
# Append the function call output as a tool message.
return ChatCompletionToolMessageParam(
diff --git a/vllm/env_override.py b/vllm/env_override.py
index 14dae2850c354..9ae1af3af46cf 100644
--- a/vllm/env_override.py
+++ b/vllm/env_override.py
@@ -95,7 +95,7 @@ def memory_plan_reuse_patched(self):
# ===================================================
# This change monkeypatches get_graph_partition_signature in pytorch 2.9.0 to
# fix inductor partition + attention-nvfp4 quant fusion, tested in
-# `tests/compile/test_fusions_e2e.py::test_attn_quant`.
+# `tests/compile/distributed/test_fusions_e2e.py::test_attn_quant`.
# For more context, see https://github.com/pytorch/pytorch/pull/165815.
diff --git a/vllm/envs.py b/vllm/envs.py
index 6bf05803e14ef..56558548d3981 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -2,8 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import functools
-import hashlib
import json
+import logging
import os
import sys
import tempfile
@@ -42,6 +42,8 @@ if TYPE_CHECKING:
VLLM_LOGGING_PREFIX: str = ""
VLLM_LOGGING_STREAM: str = "ext://sys.stdout"
VLLM_LOGGING_CONFIG_PATH: str | None = None
+ VLLM_LOGGING_COLOR: str = "auto"
+ NO_COLOR: bool = False
VLLM_LOG_STATS_INTERVAL: float = 10.0
VLLM_TRACE_FUNCTION: int = 0
VLLM_ATTENTION_BACKEND: str | None = None
@@ -50,11 +52,10 @@ if TYPE_CHECKING:
VLLM_CPU_KVCACHE_SPACE: int | None = 0
VLLM_CPU_OMP_THREADS_BIND: str = ""
VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
- VLLM_CPU_MOE_PREPACK: bool = True
VLLM_CPU_SGL_KERNEL: bool = False
VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
VLLM_XLA_CHECK_RECOMPILATION: bool = False
- VLLM_FUSED_MOE_CHUNK_SIZE: int = 64 * 1024
+ VLLM_FUSED_MOE_CHUNK_SIZE: int = 16 * 1024
VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING: bool = True
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
@@ -91,11 +92,14 @@ if TYPE_CHECKING:
VLLM_TORCH_PROFILER_DIR: str | None = None
VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False
VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
+ VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: bool = False
VLLM_USE_AOT_COMPILE: bool = False
VLLM_USE_BYTECODE_HOOK: bool = False
VLLM_FORCE_AOT_LOAD: bool = False
VLLM_TORCH_PROFILER_WITH_STACK: bool = True
VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
+ VLLM_PROFILER_DELAY_ITERS: int = 0
+ VLLM_PROFILER_MAX_ITERS: int = 0
VLLM_USE_TRITON_AWQ: bool = False
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
VLLM_SKIP_P2P_CHECK: bool = False
@@ -158,7 +162,9 @@ if TYPE_CHECKING:
VLLM_USE_FLASHINFER_MOE_FP16: bool = False
VLLM_USE_FLASHINFER_MOE_FP8: bool = False
VLLM_USE_FLASHINFER_MOE_FP4: bool = False
- VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency"] = "latency"
+ VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = (
+ "latency"
+ )
VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
VLLM_XGRAMMAR_CACHE_MB: int = 0
VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
@@ -225,7 +231,7 @@ if TYPE_CHECKING:
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
- VLLM_FLAT_LOGPROBS: bool = False
+ VLLM_USE_V2_MODEL_RUNNER: bool = False
def get_default_cache_root():
@@ -428,6 +434,8 @@ def get_vllm_port() -> int | None:
# --8<-- [start:env-vars-definition]
+logger = logging.getLogger(__name__)
+
environment_variables: dict[str, Callable[[], Any]] = {
# ================== Installation Time Env Vars ==================
# Target device of vLLM, supporting [cuda (by default),
@@ -614,6 +622,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
"VLLM_LOGGING_STREAM": lambda: os.getenv("VLLM_LOGGING_STREAM", "ext://sys.stdout"),
# if set, VLLM_LOGGING_PREFIX will be prepended to all log messages
"VLLM_LOGGING_PREFIX": lambda: os.getenv("VLLM_LOGGING_PREFIX", ""),
+ # Controls colored logging output. Options: "auto" (default, colors when terminal),
+ # "1" (always use colors), "0" (never use colors)
+ "VLLM_LOGGING_COLOR": lambda: os.getenv("VLLM_LOGGING_COLOR", "auto"),
+ # Standard unix flag for disabling ANSI color codes
+ "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0",
# If set, vllm will log stats at this interval in seconds
# If not set, vllm will log stats every 10 seconds.
"VLLM_LOG_STATS_INTERVAL": lambda: val
@@ -627,7 +640,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
# Example options:
# - "TORCH_SDPA": use torch.nn.MultiheadAttention
# - "FLASH_ATTN": use FlashAttention
- # - "XFORMERS": use XFormers
# - "FLASHINFER": use flashinfer
# - "FLASHMLA": use FlashMLA
# - "FLASH_ATTN_MLA": use FlashAttention for MLA
@@ -666,10 +678,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
)
if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
else None,
- # (CPU backend only) whether to use prepack for MoE layer. This will be
- # passed to ipex.llm.modules.GatedMLPMOE. On unsupported CPUs, you might
- # need to set this to "0" (False).
- "VLLM_CPU_MOE_PREPACK": lambda: bool(int(os.getenv("VLLM_CPU_MOE_PREPACK", "1"))),
# (CPU backend only) whether to use SGL kernels, optimized for small batch.
"VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
# If the env var is set, Ray Compiled Graph uses the specified
@@ -776,7 +784,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
# Enable SPMD mode for TPU backend.
"VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
"VLLM_FUSED_MOE_CHUNK_SIZE": lambda: int(
- os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768")
+ os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(16 * 1024))
),
# Control whether to use fused MoE activation chunking. Current chunking
# logic is incompatible with torch.compile and causes IMA. See issue
@@ -867,6 +875,19 @@ environment_variables: dict[str, Callable[[], Any]] = {
"VLLM_TORCH_PROFILER_WITH_FLOPS": lambda: bool(
os.getenv("VLLM_TORCH_PROFILER_WITH_FLOPS", "0") != "0"
),
+ # Disable torch profiling of the AsyncLLMEngine process.
+ # If set to 1, will not profile the engine process.
+ "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM": lambda: bool(
+ os.getenv("VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM", "0") != "0"
+ ),
+ # Delay number of iterations before starting profiling when using
+ # the torch/torch CUDA profiler. If set to 0, will start profiling immediately.
+ "VLLM_PROFILER_DELAY_ITERS": lambda: int(
+ os.getenv("VLLM_PROFILER_DELAY_ITERS", "0")
+ ),
+ # Maximum number of iterations to profile when using the torch/torch CUDA profiler.
+ # If set to 0, will not limit the number of iterations.
+ "VLLM_PROFILER_MAX_ITERS": lambda: int(os.getenv("VLLM_PROFILER_MAX_ITERS", "0")),
# If set, vLLM will use Triton implementations of AWQ.
"VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))),
# If set, allow loading or unloading lora adapters in runtime,
@@ -1242,7 +1263,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
# - "latency":
# Uses TensorRT-LLM kernels optimized for low-latency inference.
"VLLM_FLASHINFER_MOE_BACKEND": env_with_choices(
- "VLLM_FLASHINFER_MOE_BACKEND", "latency", ["throughput", "latency"]
+ "VLLM_FLASHINFER_MOE_BACKEND",
+ "latency",
+ ["throughput", "latency", "masked_gemm"],
),
# Control the workspace buffer size for the FlashInfer backend.
"VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int(
@@ -1267,7 +1290,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
# MoE routing strategy selector.
# See `RoutingSimulator.get_available_strategies()` # for available
# strategies.
- # Cutstom routing strategies can be registered by
+ # Custom routing strategies can be registered by
# RoutingSimulator.register_strategy()
# Note: custom strategies may not produce correct model outputs
"VLLM_MOE_ROUTING_SIMULATION_STRATEGY": lambda: os.environ.get(
@@ -1499,11 +1522,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
"VLLM_COMPILE_CACHE_SAVE_FORMAT": env_with_choices(
"VLLM_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"]
),
- # Flag to enable FlatLogprobs whose GC overhead is significantly smaller than
- # the original list[dict[int, Logprob]] approach.
- # After enabled, PromptLogprobs and SampleLogprobs would populated as
- # FlatLogprobs.
- "VLLM_FLAT_LOGPROBS": lambda: bool(int(os.getenv("VLLM_FLAT_LOGPROBS", "0"))),
+ # Flag to enable v2 model runner.
+ "VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
+ int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
+ ),
}
# --8<-- [end:env-vars-definition]
@@ -1551,85 +1573,90 @@ def is_set(name: str):
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-def compute_hash() -> str:
- """
- WARNING: Whenever a new key is added to this environment
- variables, ensure that it is included in the factors list if
- it affects the computation graph. For example, different values
- of VLLM_PP_LAYER_PARTITION will generate different computation
- graphs, so it is included in the factors list. The env vars that
- affect the choice of different kernels or attention backends should
- also be included in the factors list.
- """
+def compile_factors() -> dict[str, object]:
+ """Return env vars used for torch.compile cache keys.
- # The values of envs may affects the computation graph.
- # TODO(DefTruth): hash all environment variables?
- # for key in environment_variables:
- # factorize(key)
- environment_variables_to_hash = [
- "VLLM_PP_LAYER_PARTITION",
- "VLLM_MLA_DISABLE",
- "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH",
- "VLLM_USE_TRITON_AWQ",
- "VLLM_DP_RANK",
- "VLLM_DP_SIZE",
- "VLLM_USE_STANDALONE_COMPILE",
- "VLLM_FUSED_MOE_CHUNK_SIZE",
- "VLLM_FLASHINFER_MOE_BACKEND",
- "VLLM_V1_USE_PREFILL_DECODE_ATTENTION",
- "VLLM_ATTENTION_BACKEND",
- "VLLM_USE_FLASHINFER_SAMPLER",
- "VLLM_DISABLED_KERNELS",
- "VLLM_USE_DEEP_GEMM",
- "VLLM_MOE_USE_DEEP_GEMM",
- "VLLM_USE_DEEP_GEMM_E8M0",
- "VLLM_USE_FUSED_MOE_GROUPED_TOPK",
- "VLLM_USE_FLASHINFER_MOE_FP16",
- "VLLM_USE_FLASHINFER_MOE_FP8",
- "VLLM_USE_FLASHINFER_MOE_FP4",
- "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8",
- "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS",
- "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16",
- "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE",
- "VLLM_USE_CUDNN_PREFILL",
- "VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL",
- "VLLM_USE_TRTLLM_ATTENTION",
- "VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION",
- "VLLM_ROCM_USE_AITER",
- "VLLM_ROCM_USE_AITER_PAGED_ATTN",
- "VLLM_ROCM_USE_AITER_LINEAR",
- "VLLM_ROCM_USE_AITER_MOE",
- "VLLM_ROCM_USE_AITER_RMSNORM",
- "VLLM_ROCM_USE_AITER_MLA",
- "VLLM_ROCM_USE_AITER_MHA",
- "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
- "VLLM_ROCM_USE_AITER_TRITON_ROPE",
- "VLLM_ROCM_USE_AITER_FP8BMM",
- "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION",
- "VLLM_ROCM_USE_AITER_TRITON_GEMM",
- "VLLM_ROCM_USE_SKINNY_GEMM",
- "VLLM_ROCM_FP8_PADDING",
- "VLLM_ROCM_MOE_PADDING",
- "VLLM_ROCM_CUSTOM_PAGED_ATTN",
- "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION",
- "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16",
- "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB",
- "VLLM_ROCM_FP8_MFMA_PAGE_ATTN",
- "VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE",
- "VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING",
- "VLLM_NVFP4_GEMM_BACKEND",
- "VLLM_USE_FBGEMM",
- "VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE",
- "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL",
- ]
- for key in environment_variables_to_hash:
- # if this goes out of sync with environment_variables,
- # it's not a user error, it's a bug
- assert key in environment_variables, (
- "Please update environment_variables_to_hash in envs.py"
- )
+ Start with every known vLLM env var; drop entries in `ignored_factors`;
+ hash everything else. This keeps the cache key aligned across workers."""
- factors = [environment_variables[key]() for key in environment_variables_to_hash]
+ ignored_factors: set[str] = {
+ "MAX_JOBS",
+ "VLLM_RPC_BASE_PATH",
+ "VLLM_USE_MODELSCOPE",
+ "VLLM_RINGBUFFER_WARNING_INTERVAL",
+ "VLLM_DEBUG_DUMP_PATH",
+ "VLLM_PORT",
+ "VLLM_CACHE_ROOT",
+ "LD_LIBRARY_PATH",
+ "VLLM_SERVER_DEV_MODE",
+ "VLLM_DP_MASTER_IP",
+ "VLLM_DP_MASTER_PORT",
+ "VLLM_RANDOMIZE_DP_DUMMY_INPUTS",
+ "VLLM_CI_USE_S3",
+ "VLLM_MODEL_REDIRECT_PATH",
+ "VLLM_HOST_IP",
+ "S3_ACCESS_KEY_ID",
+ "S3_SECRET_ACCESS_KEY",
+ "S3_ENDPOINT_URL",
+ "VLLM_USAGE_STATS_SERVER",
+ "VLLM_NO_USAGE_STATS",
+ "VLLM_DO_NOT_TRACK",
+ "VLLM_LOGGING_LEVEL",
+ "VLLM_LOGGING_PREFIX",
+ "VLLM_LOGGING_STREAM",
+ "VLLM_LOGGING_CONFIG_PATH",
+ "VLLM_LOGGING_COLOR",
+ "VLLM_LOG_STATS_INTERVAL",
+ "VLLM_DEBUG_LOG_API_SERVER_RESPONSE",
+ "VLLM_TUNED_CONFIG_FOLDER",
+ "VLLM_ENGINE_ITERATION_TIMEOUT_S",
+ "VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
+ "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS",
+ "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH",
+ "VLLM_SLEEP_WHEN_IDLE",
+ "VLLM_IMAGE_FETCH_TIMEOUT",
+ "VLLM_VIDEO_FETCH_TIMEOUT",
+ "VLLM_AUDIO_FETCH_TIMEOUT",
+ "VLLM_MEDIA_URL_ALLOW_REDIRECTS",
+ "VLLM_MEDIA_LOADING_THREAD_COUNT",
+ "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB",
+ "VLLM_VIDEO_LOADER_BACKEND",
+ "VLLM_MEDIA_CONNECTOR",
+ "VLLM_ASSETS_CACHE",
+ "VLLM_ASSETS_CACHE_MODEL_CLEAN",
+ "VLLM_MM_INPUT_CACHE_GIB",
+ "VLLM_WORKER_MULTIPROC_METHOD",
+ "VLLM_ENABLE_V1_MULTIPROCESSING",
+ "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE",
+ "VLLM_CPU_KVCACHE_SPACE",
+ "VLLM_CPU_OMP_THREADS_BIND",
+ "VLLM_CPU_NUM_OF_RESERVED_CPU",
+ "VLLM_CPU_MOE_PREPACK",
+ "VLLM_CPU_SGL_KERNEL",
+ "VLLM_TEST_FORCE_LOAD_FORMAT",
+ "LOCAL_RANK",
+ "CUDA_VISIBLE_DEVICES",
+ "NO_COLOR",
+ }
+
+ from vllm.config.utils import normalize_value
+
+ factors: dict[str, object] = {}
+ for factor, getter in environment_variables.items():
+ if factor in ignored_factors:
+ continue
+
+ try:
+ raw = getter()
+ except Exception as exc: # pragma: no cover - defensive logging
+ logger.warning(
+ "Skipping environment variable %s while hashing compile factors: %s",
+ factor,
+ exc,
+ )
+ continue
+
+ factors[factor] = normalize_value(raw)
ray_noset_env_vars = [
# Refer to
@@ -1652,8 +1679,8 @@ def compute_hash() -> str:
"RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR",
"RAY_EXPERIMENTAL_NOSET_RBLN_RT_VISIBLE_DEVICES",
]
- factors.extend([os.getenv(var) for var in ray_noset_env_vars])
- hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
+ for var in ray_noset_env_vars:
+ factors[var] = normalize_value(os.getenv(var))
- return hash_str
+ return factors
diff --git a/vllm/logger.py b/vllm/logger.py
index 9341008296843..772e36497b45e 100644
--- a/vllm/logger.py
+++ b/vllm/logger.py
@@ -17,18 +17,25 @@ from typing import Any, Literal, cast
import vllm.envs as envs
-VLLM_CONFIGURE_LOGGING = envs.VLLM_CONFIGURE_LOGGING
-VLLM_LOGGING_CONFIG_PATH = envs.VLLM_LOGGING_CONFIG_PATH
-VLLM_LOGGING_LEVEL = envs.VLLM_LOGGING_LEVEL
-VLLM_LOGGING_PREFIX = envs.VLLM_LOGGING_PREFIX
-VLLM_LOGGING_STREAM = envs.VLLM_LOGGING_STREAM
-
_FORMAT = (
- f"{VLLM_LOGGING_PREFIX}%(levelname)s %(asctime)s "
+ f"{envs.VLLM_LOGGING_PREFIX}%(levelname)s %(asctime)s "
"[%(fileinfo)s:%(lineno)d] %(message)s"
)
_DATE_FORMAT = "%m-%d %H:%M:%S"
+
+def _use_color() -> bool:
+ if envs.NO_COLOR or envs.VLLM_LOGGING_COLOR == "0":
+ return False
+ if envs.VLLM_LOGGING_COLOR == "1":
+ return True
+ if envs.VLLM_LOGGING_STREAM == "ext://sys.stdout": # stdout
+ return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()
+ elif envs.VLLM_LOGGING_STREAM == "ext://sys.stderr": # stderr
+ return hasattr(sys.stderr, "isatty") and sys.stderr.isatty()
+ return False
+
+
DEFAULT_LOGGING_CONFIG = {
"formatters": {
"vllm": {
@@ -36,13 +43,19 @@ DEFAULT_LOGGING_CONFIG = {
"datefmt": _DATE_FORMAT,
"format": _FORMAT,
},
+ "vllm_color": {
+ "class": "vllm.logging_utils.ColoredFormatter",
+ "datefmt": _DATE_FORMAT,
+ "format": _FORMAT,
+ },
},
"handlers": {
"vllm": {
"class": "logging.StreamHandler",
- "formatter": "vllm",
- "level": VLLM_LOGGING_LEVEL,
- "stream": VLLM_LOGGING_STREAM,
+ # Choose formatter based on color setting.
+ "formatter": "vllm_color" if _use_color() else "vllm",
+ "level": envs.VLLM_LOGGING_LEVEL,
+ "stream": envs.VLLM_LOGGING_STREAM,
},
},
"loggers": {
@@ -144,7 +157,7 @@ _METHODS_TO_PATCH = {
def _configure_vllm_root_logger() -> None:
logging_config = dict[str, Any]()
- if not VLLM_CONFIGURE_LOGGING and VLLM_LOGGING_CONFIG_PATH:
+ if not envs.VLLM_CONFIGURE_LOGGING and envs.VLLM_LOGGING_CONFIG_PATH:
raise RuntimeError(
"VLLM_CONFIGURE_LOGGING evaluated to false, but "
"VLLM_LOGGING_CONFIG_PATH was given. VLLM_LOGGING_CONFIG_PATH "
@@ -152,16 +165,22 @@ def _configure_vllm_root_logger() -> None:
"VLLM_CONFIGURE_LOGGING or unset VLLM_LOGGING_CONFIG_PATH."
)
- if VLLM_CONFIGURE_LOGGING:
+ if envs.VLLM_CONFIGURE_LOGGING:
logging_config = DEFAULT_LOGGING_CONFIG
- if VLLM_LOGGING_CONFIG_PATH:
- if not path.exists(VLLM_LOGGING_CONFIG_PATH):
+ vllm_handler = logging_config["handlers"]["vllm"]
+ # Refresh these values in case env vars have changed.
+ vllm_handler["level"] = envs.VLLM_LOGGING_LEVEL
+ vllm_handler["stream"] = envs.VLLM_LOGGING_STREAM
+ vllm_handler["formatter"] = "vllm_color" if _use_color() else "vllm"
+
+ if envs.VLLM_LOGGING_CONFIG_PATH:
+ if not path.exists(envs.VLLM_LOGGING_CONFIG_PATH):
raise RuntimeError(
"Could not load logging config. File does not exist: %s",
- VLLM_LOGGING_CONFIG_PATH,
+ envs.VLLM_LOGGING_CONFIG_PATH,
)
- with open(VLLM_LOGGING_CONFIG_PATH, encoding="utf-8") as file:
+ with open(envs.VLLM_LOGGING_CONFIG_PATH, encoding="utf-8") as file:
custom_config = json.loads(file.read())
if not isinstance(custom_config, dict):
diff --git a/vllm/logging_utils/__init__.py b/vllm/logging_utils/__init__.py
index 7202259ca21aa..8d3354df215b1 100644
--- a/vllm/logging_utils/__init__.py
+++ b/vllm/logging_utils/__init__.py
@@ -1,10 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from vllm.logging_utils.formatter import NewLineFormatter
+from vllm.logging_utils.formatter import ColoredFormatter, NewLineFormatter
+from vllm.logging_utils.lazy import lazy
from vllm.logging_utils.log_time import logtime
__all__ = [
"NewLineFormatter",
+ "ColoredFormatter",
+ "lazy",
"logtime",
]
diff --git a/vllm/logging_utils/formatter.py b/vllm/logging_utils/formatter.py
index 02ba308e18796..3ad4ef8d119ad 100644
--- a/vllm/logging_utils/formatter.py
+++ b/vllm/logging_utils/formatter.py
@@ -75,3 +75,53 @@ class NewLineFormatter(logging.Formatter):
parts = msg.split(record.message)
msg = msg.replace("\n", "\r\n" + parts[0])
return msg
+
+
+class ColoredFormatter(NewLineFormatter):
+ """Adds ANSI color codes to log levels for terminal output.
+
+ This formatter adds colors by injecting them into the format string for
+ static elements (timestamp, filename, line number) and modifying the
+ levelname attribute for dynamic color selection.
+ """
+
+ # ANSI color codes
+ COLORS = {
+ "DEBUG": "\033[37m", # White
+ "INFO": "\033[32m", # Green
+ "WARNING": "\033[33m", # Yellow
+ "ERROR": "\033[31m", # Red
+ "CRITICAL": "\033[35m", # Magenta
+ }
+ GREY = "\033[90m" # Grey for timestamp and file info
+ RESET = "\033[0m"
+
+ def __init__(self, fmt, datefmt=None, style="%"):
+ # Inject grey color codes into format string for timestamp and file info
+ if fmt:
+ # Wrap %(asctime)s with grey
+ fmt = fmt.replace("%(asctime)s", f"{self.GREY}%(asctime)s{self.RESET}")
+ # Wrap [%(fileinfo)s:%(lineno)d] with grey
+ fmt = fmt.replace(
+ "[%(fileinfo)s:%(lineno)d]",
+ f"{self.GREY}[%(fileinfo)s:%(lineno)d]{self.RESET}",
+ )
+
+ # Call parent __init__ with potentially modified format string
+ super().__init__(fmt, datefmt, style)
+
+ def format(self, record):
+ # Store original levelname to restore later (in case record is reused)
+ orig_levelname = record.levelname
+
+ # Only modify levelname - it needs dynamic color based on severity
+ if (color_code := self.COLORS.get(record.levelname)) is not None:
+ record.levelname = f"{color_code}{record.levelname}{self.RESET}"
+
+ # Call parent format which will handle everything else
+ msg = super().format(record)
+
+ # Restore original levelname
+ record.levelname = orig_levelname
+
+ return msg
diff --git a/vllm/logging_utils/lazy.py b/vllm/logging_utils/lazy.py
new file mode 100644
index 0000000000000..3ade798962857
--- /dev/null
+++ b/vllm/logging_utils/lazy.py
@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from typing import Any
+
+
+class lazy:
+ """Wrap a zero-argument callable evaluated only during log formatting."""
+
+ __slots__ = ("_factory",)
+
+ def __init__(self, factory: Callable[[], Any]) -> None:
+ self._factory = factory
+
+ def __str__(self) -> str:
+ return str(self._factory())
+
+ def __repr__(self) -> str:
+ return str(self)
diff --git a/vllm/logprobs.py b/vllm/logprobs.py
index a34398db2c960..6a820308f523f 100644
--- a/vllm/logprobs.py
+++ b/vllm/logprobs.py
@@ -5,8 +5,6 @@ from collections.abc import Iterable, Iterator, MutableSequence
from dataclasses import dataclass, field
from typing import overload
-import vllm.envs as envs
-
# We use dataclass for now because it is used for
# openai server output, and msgspec is not serializable.
@@ -161,17 +159,17 @@ PromptLogprobs = FlatLogprobs | list[LogprobsOnePosition | None]
SampleLogprobs = FlatLogprobs | list[LogprobsOnePosition]
-def create_prompt_logprobs() -> PromptLogprobs:
+def create_prompt_logprobs(flat_logprobs: bool) -> PromptLogprobs:
"""Creates a container to store prompt logprobs for a request"""
- logprobs = FlatLogprobs() if envs.VLLM_FLAT_LOGPROBS else []
+ logprobs = FlatLogprobs() if flat_logprobs else []
# NOTE: logprob of first prompt token is None.
logprobs.append(None)
return logprobs
-def create_sample_logprobs() -> SampleLogprobs:
+def create_sample_logprobs(flat_logprobs: bool) -> SampleLogprobs:
"""Creates a container to store decode logprobs for a request"""
- return FlatLogprobs() if envs.VLLM_FLAT_LOGPROBS else []
+ return FlatLogprobs() if flat_logprobs else []
def append_logprobs_for_next_position(
diff --git a/vllm/lora/layers/__init__.py b/vllm/lora/layers/__init__.py
index 8a4f5ff175d4f..25364a5881364 100644
--- a/vllm/lora/layers/__init__.py
+++ b/vllm/lora/layers/__init__.py
@@ -11,7 +11,7 @@ from vllm.lora.layers.column_parallel_linear import (
QKVParallelLinearWithLoRA,
QKVParallelLinearWithShardedLoRA,
)
-from vllm.lora.layers.fused_moe import FusedMoEWithLoRA
+from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA
from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA
from vllm.lora.layers.row_parallel_linear import (
@@ -38,4 +38,5 @@ __all__ = [
"ReplicatedLinearWithLoRA",
"LoRAMapping",
"FusedMoEWithLoRA",
+ "FusedMoE3DWithLoRA",
]
diff --git a/vllm/lora/layers/base.py b/vllm/lora/layers/base.py
index 0c7e806848892..3bfb88c007622 100644
--- a/vllm/lora/layers/base.py
+++ b/vllm/lora/layers/base.py
@@ -42,9 +42,8 @@ class BaseLayerWithLoRA(nn.Module):
def set_lora(
self,
index: int,
- lora_a: torch.Tensor,
- lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
):
"""Overwrites lora tensors at index."""
...
diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py
index 3db4165e20176..06ecc8d2f634c 100644
--- a/vllm/lora/layers/base_linear.py
+++ b/vllm/lora/layers/base_linear.py
@@ -94,14 +94,15 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
def set_lora(
self,
index: int,
- lora_a: torch.Tensor,
- lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
):
# Except for QKVParallelLinearWithLoRA and
# MergedColumnParallelLinearWithLoRA, all other linear LoRA layers
# store weights in a tuple of size 1. These two layers will
# override this function.
+ assert isinstance(lora_a, torch.Tensor)
+ assert isinstance(lora_b, torch.Tensor)
assert (
len(self.lora_a_stacked) == len(self.lora_b_stacked) == self.n_slices == 1
)
diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py
index 637ded9b2a0f0..3e21d426c304a 100644
--- a/vllm/lora/layers/column_parallel_linear.py
+++ b/vllm/lora/layers/column_parallel_linear.py
@@ -246,9 +246,8 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
def set_lora(
self,
index: int,
- lora_a: torch.Tensor,
- lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
):
self.reset_lora(index)
diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py
index 8fb3efa220f6d..0eb6562bec6cd 100644
--- a/vllm/lora/layers/fused_moe.py
+++ b/vllm/lora/layers/fused_moe.py
@@ -12,6 +12,7 @@ from vllm.distributed.parallel_state import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
+from vllm.distributed.utils import divide
from vllm.lora.layers.base import BaseLayerWithLoRA
from vllm.lora.ops.triton_ops.utils import get_lora_op_configs
from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -41,6 +42,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
self.tp_size = get_tensor_model_parallel_world_size()
self.tp_rank = get_tensor_model_parallel_rank()
self.device = base_layer.w2_weight.device
+ self._w13_slices = 2
self._inject_lora_into_fused_moe()
def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
@@ -59,8 +61,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
def _get_lora_moe_configs(
self,
op_prefix: str,
- lora_a_stacked: torch.Tensor,
- lora_b_stacked: torch.Tensor,
+ num_loras: int,
+ rank: int,
num_slices: int,
M: int,
layer: FusedMoE,
@@ -68,23 +70,25 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
config_dtype: str,
):
if envs.VLLM_TUNED_CONFIG_FOLDER:
+ hidden_size = layer.hidden_size
+ intermediate_size = layer.intermediate_size_per_partition
shrink_config = get_lora_op_configs(
op_type=f"fused_moe_lora_{op_prefix}_shrink",
- max_loras=lora_a_stacked.shape[0],
+ max_loras=num_loras,
batch=M,
- hidden_size=lora_a_stacked.shape[-1],
- rank=lora_a_stacked.shape[-2],
+ hidden_size=hidden_size,
+ rank=rank,
num_slices=num_slices,
- moe_intermediate_size=lora_b_stacked.shape[-2],
+ moe_intermediate_size=intermediate_size,
)
expand_config = get_lora_op_configs(
op_type=f"fused_moe_lora_{op_prefix}_expand",
- max_loras=lora_a_stacked.shape[0],
+ max_loras=num_loras,
batch=M,
- hidden_size=lora_a_stacked.shape[-1],
- rank=lora_a_stacked.shape[-2],
+ hidden_size=hidden_size, # lora_a_stacked.shape[-1],
+ rank=rank,
num_slices=num_slices,
- moe_intermediate_size=lora_b_stacked.shape[-2],
+ moe_intermediate_size=intermediate_size, # lora_b_stacked.shape[-2],
)
else: # fall back to the default config
get_config_func = functools.partial(
@@ -151,12 +155,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
num_tokens = hidden_states.size(0)
M = min(num_tokens, CHUNK_SIZE)
-
+ max_lora_rank = self.w13_lora_a_stacked[0].shape[-2]
shrink_config, expand_config = self._get_lora_moe_configs(
op_prefix="w13",
- lora_a_stacked=self.w1_lora_a_stacked,
- lora_b_stacked=self.w1_lora_b_stacked,
- num_slices=2,
+ num_loras=self.max_loras,
+ rank=max_lora_rank,
+ num_slices=self._w13_slices,
M=M,
layer=layer,
top_k=top_k,
@@ -164,7 +168,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
)
# get the block size of m from customized config or default config
- max_loras = self.w1_lora_a_stacked.shape[0]
(
sorted_token_ids_lora,
expert_ids_lora,
@@ -174,7 +177,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
num_tokens,
shrink_config["BLOCK_SIZE_M"],
self.base_layer.local_num_experts,
- max_loras,
+ self.max_loras,
self.adapter_enabled,
expert_map,
)
@@ -185,17 +188,15 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
num_tokens_post_padded_lora
)
- w13_lora_a_stacked = [self.w1_lora_a_stacked, self.w3_lora_a_stacked]
- w13_lora_b_stacked = [self.w1_lora_b_stacked, self.w3_lora_b_stacked]
- max_lora_rank = self.w1_lora_a_stacked.shape[-2]
- expert_ids_lora = expert_ids_lora.view(max_loras, -1)
- sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
+ expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
+ sorted_token_ids_lora = sorted_token_ids_lora.view(self.max_loras, -1)
+ #
self.punica_wrapper.add_lora_fused_moe(
input.view(-1, top_k, input.shape[-1]),
hidden_states,
- w13_lora_a_stacked,
- w13_lora_b_stacked,
+ self.w13_lora_a_stacked,
+ self.w13_lora_b_stacked,
topk_weights,
sorted_token_ids_lora,
expert_ids_lora,
@@ -205,6 +206,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
shrink_config, ## pass the shrink config
expand_config, ## pass the expand config
self.adapter_enabled,
+ fully_sharded=self.fully_sharded,
)
result = func(*args, **kwargs)
@@ -228,11 +230,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
num_tokens = hidden_states.size(0)
M = min(num_tokens, CHUNK_SIZE)
-
+ max_lora_rank = self.w2_lora_a_stacked[0].shape[-2]
shrink_config, expand_config = self._get_lora_moe_configs(
op_prefix="w2",
- lora_a_stacked=self.w2_lora_a_stacked,
- lora_b_stacked=self.w2_lora_b_stacked,
+ num_loras=self.max_loras,
+ rank=max_lora_rank,
num_slices=1,
M=M,
layer=layer,
@@ -245,17 +247,19 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
num_tokens_post_padded_lora = moe_state_dict[
"num_tokens_post_padded_lora"
]
- max_loras = self.w1_lora_a_stacked.shape[0]
- expert_ids_lora = expert_ids_lora.view(max_loras, -1)
- sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
+
+ expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
+ sorted_token_ids_lora = sorted_token_ids_lora.view(self.max_loras, -1)
intermediate_cache2 = moe_state_dict["intermediate_cache2"]
intermediate_cache3 = args[0]
- max_lora_rank = self.w1_lora_a_stacked.shape[-2]
+
+ shard_size_w2 = divide(self.base_layer.hidden_size, self.tp_size)
+
self.punica_wrapper.add_lora_fused_moe(
intermediate_cache3,
intermediate_cache2,
- [self.w2_lora_a_stacked],
- [self.w2_lora_b_stacked],
+ self.w2_lora_a_stacked,
+ self.w2_lora_b_stacked,
topk_weights,
sorted_token_ids_lora,
expert_ids_lora,
@@ -266,6 +270,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
expand_config, ## pass the expand config
self.adapter_enabled,
True,
+ fully_sharded=self.fully_sharded,
+ offset=shard_size_w2 * self.tp_rank if self.fully_sharded else 0,
)
result = func(*args, **kwargs)
@@ -282,11 +288,72 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
fused_experts.moe_sum = moe_sum_decorator(
self.base_layer, fused_experts.moe_sum
)
-
self.base_layer.quant_method = FusedMoEModularMethod(
self.base_layer.quant_method, m_fused_moe_fn
)
+ def _create_lora_a_weights(
+ self,
+ max_loras: int,
+ lora_config: LoRAConfig,
+ ):
+ self.w13_lora_a_stacked: tuple[torch.Tensor, ...] = tuple(
+ torch.zeros(
+ (
+ max_loras,
+ self.base_layer.local_num_experts,
+ lora_config.max_lora_rank
+ if not self.fully_sharded
+ else divide(lora_config.max_lora_rank, self.tp_size),
+ self.base_layer.hidden_size,
+ ),
+ dtype=lora_config.lora_dtype,
+ device=self.device,
+ )
+ for _ in range(self._w13_slices)
+ )
+ self.w2_lora_a_stacked: tuple[torch.Tensor, ...] = (
+ torch.zeros(
+ (
+ max_loras,
+ self.base_layer.local_num_experts,
+ lora_config.max_lora_rank,
+ self.base_layer.intermediate_size_per_partition,
+ ),
+ dtype=lora_config.lora_dtype,
+ device=self.device,
+ ),
+ )
+
+ def _create_lora_b_weights(self, max_loras: int, lora_config: LoRAConfig):
+ self.w13_lora_b_stacked: tuple[torch.Tensor, ...] = tuple(
+ torch.zeros(
+ (
+ max_loras,
+ self.base_layer.local_num_experts,
+ self.base_layer.intermediate_size_per_partition,
+ lora_config.max_lora_rank,
+ ),
+ dtype=lora_config.lora_dtype,
+ device=self.device,
+ )
+ for _ in range(self._w13_slices)
+ )
+ self.w2_lora_b_stacked: tuple[torch.Tensor, ...] = (
+ torch.zeros(
+ (
+ max_loras,
+ self.base_layer.local_num_experts,
+ self.base_layer.hidden_size
+ if not self.fully_sharded
+ else divide(self.base_layer.hidden_size, self.tp_size),
+ lora_config.max_lora_rank,
+ ),
+ dtype=lora_config.lora_dtype,
+ device=self.device,
+ ),
+ )
+
def create_lora_weights(
self,
max_loras: int,
@@ -294,108 +361,63 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
model_config: PretrainedConfig | None = None,
) -> None:
"""Initializes lora matrices."""
+ self.max_loras = lora_config.max_loras
+ self.fully_sharded = lora_config.fully_sharded_loras
self.adapter_enabled = torch.tensor(
[0] * (max_loras + 1), dtype=torch.int, device=self.device
)
- self.w1_lora_a_stacked = torch.zeros(
- (
- max_loras,
- self.base_layer.local_num_experts,
- lora_config.max_lora_rank,
- self.base_layer.hidden_size,
- ),
- dtype=lora_config.lora_dtype,
- device=self.device,
- )
- self.w1_lora_b_stacked = torch.zeros(
- (
- max_loras,
- self.base_layer.local_num_experts,
- self.base_layer.intermediate_size_per_partition,
- lora_config.max_lora_rank,
- ),
- dtype=lora_config.lora_dtype,
- device=self.device,
- )
-
- self.w2_lora_a_stacked = torch.zeros(
- (
- max_loras,
- self.base_layer.local_num_experts,
- lora_config.max_lora_rank,
- self.base_layer.intermediate_size_per_partition,
- ),
- dtype=lora_config.lora_dtype,
- device=self.device,
- )
- self.w2_lora_b_stacked = torch.zeros(
- (
- max_loras,
- self.base_layer.local_num_experts,
- self.base_layer.hidden_size,
- lora_config.max_lora_rank,
- ),
- dtype=lora_config.lora_dtype,
- device=self.device,
- )
-
- self.w3_lora_a_stacked = torch.zeros(
- (
- max_loras,
- self.base_layer.local_num_experts,
- lora_config.max_lora_rank,
- self.base_layer.hidden_size,
- ),
- dtype=lora_config.lora_dtype,
- device=self.device,
- )
- self.w3_lora_b_stacked = torch.zeros(
- (
- max_loras,
- self.base_layer.local_num_experts,
- self.base_layer.intermediate_size_per_partition,
- lora_config.max_lora_rank,
- ),
- dtype=lora_config.lora_dtype,
- device=self.device,
- )
-
+ self._create_lora_a_weights(max_loras, lora_config)
+ self._create_lora_b_weights(max_loras, lora_config)
# They will be used by 'LoRALayerWeights.create_dummy_lora_weights'
# to create a dummy LoRA weights.
+ # TODO Optimize this section
self.lora_a_stacked = []
self.lora_b_stacked = []
for lora_id in range(max_loras):
for experts_id in range(self.base_layer.local_num_experts):
# gate_proj,down_proj,up_proj
- self.lora_a_stacked.append(self.w1_lora_a_stacked[lora_id][experts_id])
- self.lora_a_stacked.append(self.w2_lora_a_stacked[lora_id][experts_id])
- self.lora_a_stacked.append(self.w3_lora_a_stacked[lora_id][experts_id])
+ self.lora_a_stacked.append(
+ self.w13_lora_a_stacked[0][lora_id][experts_id]
+ )
+ self.lora_a_stacked.append(
+ self.w2_lora_a_stacked[0][lora_id][experts_id]
+ )
- self.lora_b_stacked.append(self.w1_lora_b_stacked[lora_id][experts_id])
- self.lora_b_stacked.append(self.w2_lora_b_stacked[lora_id][experts_id])
- self.lora_b_stacked.append(self.w3_lora_b_stacked[lora_id][experts_id])
+ self.lora_b_stacked.append(
+ self.w13_lora_b_stacked[0][lora_id][experts_id]
+ )
+ self.lora_b_stacked.append(
+ self.w2_lora_b_stacked[0][lora_id][experts_id]
+ )
+
+ self.lora_a_stacked.append(
+ self.w13_lora_a_stacked[1][lora_id][experts_id]
+ )
+ self.lora_b_stacked.append(
+ self.w13_lora_b_stacked[1][lora_id][experts_id]
+ )
def reset_lora(self, index: int):
"""Resets the lora weights at index back to 0."""
- self.w1_lora_a_stacked[index] = 0
- self.w1_lora_b_stacked[index] = 0
- self.w3_lora_a_stacked[index] = 0
- self.w3_lora_b_stacked[index] = 0
- self.w2_lora_a_stacked[index] = 0
- self.w2_lora_b_stacked[index] = 0
+ for pos in range(self._w13_slices):
+ self.w13_lora_a_stacked[pos][index] = 0
+ self.w13_lora_b_stacked[pos][index] = 0
+
+ self.w2_lora_a_stacked[0][index] = 0
+ self.w2_lora_b_stacked[0][index] = 0
self.adapter_enabled[index] = 0
def set_lora(
self,
index: int,
- lora_a: torch.Tensor,
- lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None,
- bias: torch.Tensor | None = None,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
):
"""Overwrites lora tensors at index."""
+ assert isinstance(lora_a, list)
+ assert isinstance(lora_b, list)
self.reset_lora(index)
self.adapter_enabled[index] = 1
for eid in range(len(lora_a) // 3):
@@ -419,39 +441,44 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
w3_lora_b = w3_lora_b[start_idx:end_idx, :]
w2_lora_a = w2_lora_a[:, start_idx:end_idx]
- self.w1_lora_a_stacked[
+ if self.fully_sharded:
+ # Based on S-LoRA, we slice W1 and W3 A along the rank dim,
+ # and W2 B along the hidden_size dim.
+ w13_shard_size = self.w13_lora_a_stacked[0][index, eid].shape[0]
+ w13_start_idx = self.tp_rank * w13_shard_size
+ w13_end_idx = (self.tp_rank + 1) * w13_shard_size
+ w1_lora_a = w1_lora_a[w13_start_idx:w13_end_idx, :]
+ w3_lora_a = w3_lora_a[w13_start_idx:w13_end_idx, :]
+
+ w2_shard_size = self.w2_lora_b_stacked[0][index, eid].shape[0]
+ w2_start_idx = self.tp_rank * w2_shard_size
+ w2_end_idx = (self.tp_rank + 1) * w2_shard_size
+ w2_lora_b = w2_lora_b[w2_start_idx:w2_end_idx, :]
+ # w1 lora_a
+ self.w13_lora_a_stacked[0][
index, eid, : w1_lora_a.shape[0], : w1_lora_a.shape[1]
].copy_(w1_lora_a, non_blocking=True)
-
- self.w3_lora_a_stacked[
+ # w3 lora_a
+ self.w13_lora_a_stacked[1][
index, eid, : w3_lora_a.shape[0], : w3_lora_a.shape[1]
].copy_(w3_lora_a, non_blocking=True)
- self.w2_lora_b_stacked[
- index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
- ].copy_(w2_lora_b, non_blocking=True)
-
- self.w1_lora_b_stacked[
+ # w1 lora_b
+ self.w13_lora_b_stacked[0][
index, eid, : w1_lora_b.shape[0], : w1_lora_b.shape[1]
].copy_(w1_lora_b, non_blocking=True)
- self.w3_lora_b_stacked[
+ # w3 lora_b
+ self.w13_lora_b_stacked[1][
index, eid, : w3_lora_b.shape[0], : w3_lora_b.shape[1]
].copy_(w3_lora_b, non_blocking=True)
- self.w2_lora_a_stacked[
+
+ self.w2_lora_a_stacked[0][
index, eid, : w2_lora_a.shape[0], : w2_lora_a.shape[1]
].copy_(w2_lora_a, non_blocking=True)
- @classmethod
- def can_replace_layer(
- cls,
- source_layer: nn.Module,
- lora_config: LoRAConfig,
- packed_modules_list: list,
- model_config: PretrainedConfig | None,
- ) -> bool:
- """Returns True if the layer can be replaced by this LoRA layer."""
- # return type(source_layer) is FusedMoE
- return isinstance(source_layer, FusedMoE)
+ self.w2_lora_b_stacked[0][
+ index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
+ ].copy_(w2_lora_b, non_blocking=True)
def forward(self, *args, **kwargs):
return self.base_layer.forward(*args, **kwargs)
@@ -470,3 +497,220 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
@property
def is_internal_router(self) -> bool:
return self.base_layer.is_internal_router
+
+ @classmethod
+ def can_replace_layer(
+ cls,
+ source_layer: nn.Module,
+ lora_config: LoRAConfig,
+ packed_modules_list: list,
+ model_config: PretrainedConfig | None,
+ ) -> bool:
+ """Returns True if the layer can be replaced by this LoRA layer."""
+ # return type(source_layer) is FusedMoE
+
+ return type(source_layer) is FusedMoE and len(packed_modules_list) == 2
+
+
+class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
+ def __init__(self, base_layer):
+ super().__init__(base_layer)
+ self._w13_slices = 1
+
+ def _create_lora_b_weights(self, max_loras, lora_config):
+ self.w13_lora_b_stacked: tuple[torch.Tensor] = tuple(
+ torch.zeros(
+ (
+ max_loras,
+ self.base_layer.local_num_experts,
+ self.base_layer.intermediate_size_per_partition * 2,
+ lora_config.max_lora_rank,
+ ),
+ dtype=lora_config.lora_dtype,
+ device=self.device,
+ )
+ for _ in range(self._w13_slices)
+ )
+ self.w2_lora_b_stacked: tuple[torch.Tensor] = (
+ torch.zeros(
+ (
+ max_loras,
+ self.base_layer.local_num_experts,
+ self.base_layer.hidden_size
+ if not self.fully_sharded
+ else divide(self.base_layer.hidden_size, self.tp_size),
+ lora_config.max_lora_rank,
+ ),
+ dtype=lora_config.lora_dtype,
+ device=self.device,
+ ),
+ )
+
+ def create_lora_weights(
+ self,
+ max_loras: int,
+ lora_config: LoRAConfig,
+ model_config: PretrainedConfig | None = None,
+ ) -> None:
+ """Initializes lora matrices."""
+ self.max_loras = lora_config.max_loras
+ self.fully_sharded = lora_config.fully_sharded_loras
+
+ self.adapter_enabled = torch.tensor(
+ [0] * (max_loras + 1), dtype=torch.int, device=self.device
+ )
+
+ self._create_lora_a_weights(max_loras, lora_config)
+ self._create_lora_b_weights(max_loras, lora_config)
+
+ def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor:
+ if self.tp_size == 1 or not self.fully_sharded:
+ return w13_lora_a
+
+ # w13_lora_a shape (num_experts,rank,input_size)
+ current_lora_rank = w13_lora_a.shape[1]
+ assert current_lora_rank % self.tp_size == 0
+
+ sliced_rank = current_lora_rank // self.tp_size
+ start_idx = self.tp_rank * sliced_rank
+ end_idx = (self.tp_rank + 1) * sliced_rank
+ return w13_lora_a[:, start_idx:end_idx, :]
+
+ def _slice_w13_b(self, w13_lora_b: torch.Tensor, is_interleave: bool = True):
+ if self.tp_size == 1:
+ return w13_lora_b
+
+ # w13_lora_b shape (num_experts,output_size,rank)
+ shard_size = self.base_layer.intermediate_size_per_partition
+ start_idx = self.tp_rank * shard_size
+ end_idx = (self.tp_rank + 1) * shard_size
+ if is_interleave:
+ # For models like GPT-OSS, the weights of w1 (gate_proj) and w3 (up_proj)
+ # in the interleaved order, and corresponding LoRA need to be processed.
+ w1_lora_b = w13_lora_b[:, ::2, :]
+ w3_lora_b = w13_lora_b[:, 1::2, :]
+ sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :]
+ sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :]
+
+ return torch.stack([sliced_w1_lora_b, sliced_w3_lora_b], dim=2).flatten(
+ 1, 2
+ )
+ else:
+ slice_size = w13_lora_b.shape[1] // 2
+ w1_lora_b = w13_lora_b[:, :slice_size, :]
+ w3_lora_b = w13_lora_b[:, slice_size:, :]
+ sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :]
+ sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :]
+
+ return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1)
+
+ def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor:
+ if self.tp_size == 1:
+ return w2_lora_a
+ # w2_lora_a shape (num_experts,rank,input_size)
+ shard_size = self.base_layer.intermediate_size_per_partition
+ start_idx = self.tp_rank * shard_size
+ end_idx = (self.tp_rank + 1) * shard_size
+
+ return w2_lora_a[:, :, start_idx:end_idx]
+
+ def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor:
+ if self.tp_size == 1 or not self.fully_sharded:
+ return w2_lora_b
+ # Based on S-LoRA, we slice W2 B along the hidden_size dim.
+ # w2_lora_b shape (num_experts,output_size,rank)
+ current_lora_size = w2_lora_b.shape[1]
+
+ sliced_size = current_lora_size // self.tp_size
+ start_idx = self.tp_rank * sliced_size
+ end_idx = (self.tp_rank + 1) * sliced_size
+ return w2_lora_b[:, start_idx:end_idx, :]
+
+ def set_lora(
+ self,
+ index: int,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
+ ):
+ """Overwrites lora tensors at index."""
+ # Make mypy happy
+ assert isinstance(lora_a, list)
+ assert isinstance(lora_b, list)
+ assert len(lora_a) == len(lora_b) == 2
+
+ self.reset_lora(index)
+ self.adapter_enabled[index] = 1
+
+ num_experts = self.w13_lora_a_stacked[0].shape[1]
+ w13_lora_a, w2_lora_a = lora_a
+ w13_lora_b, w2_lora_b = lora_b
+
+ # (num_experts,rank,input_size)
+ w13_lora_a = w13_lora_a.reshape(num_experts, -1, w13_lora_a.shape[-1])
+ w2_lora_a = w2_lora_a.reshape(num_experts, -1, w2_lora_a.shape[-1])
+ # (output_size,num_experts,rank)
+ w13_lora_b = w13_lora_b.reshape(w13_lora_b.shape[0], num_experts, -1)
+ w2_lora_b = w2_lora_b.reshape(w2_lora_b.shape[0], num_experts, -1)
+ # (num_experts,output_size,rank)
+ w13_lora_b = w13_lora_b.permute(1, 0, 2)
+ w2_lora_b = w2_lora_b.permute(1, 0, 2)
+
+ sliced_w13_lora_a = self._slice_w13_a(w13_lora_a)
+ sliced_w13_lora_b = self._slice_w13_b(w13_lora_b, is_interleave=True)
+
+ sliced_w2_lora_a = self._slice_w2_a(w2_lora_a)
+ sliced_w2_lora_b = self._slice_w2_b(w2_lora_b)
+
+ self.w13_lora_a_stacked[0][
+ index, :, : sliced_w13_lora_a.shape[1], : sliced_w13_lora_a.shape[2]
+ ].copy_(sliced_w13_lora_a, non_blocking=True)
+ self.w2_lora_a_stacked[0][
+ index, :, : sliced_w2_lora_a.shape[1], : sliced_w2_lora_a.shape[2]
+ ].copy_(sliced_w2_lora_a, non_blocking=True)
+
+ self.w13_lora_b_stacked[0][
+ index, :, : sliced_w13_lora_b.shape[1], : sliced_w13_lora_b.shape[2]
+ ].copy_(sliced_w13_lora_b, non_blocking=True)
+ self.w2_lora_b_stacked[0][
+ index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2]
+ ].copy_(sliced_w2_lora_b, non_blocking=True)
+
+ @property
+ def w13_input_size(self):
+ """
+ Full size
+ """
+ return self.w13_lora_a_stacked[0].shape[-1]
+
+ @property
+ def w13_output_size(self):
+ """
+ Full size
+ """
+ return self.w13_lora_b_stacked[0].shape[-2] * self.tp_size
+
+ @property
+ def w2_input_size(self):
+ """
+ Full size
+ """
+ return self.w2_lora_a_stacked[0].shape[-1] * self.tp_size
+
+ @property
+ def w2_output_size(self):
+ """
+ Full size
+ """
+ return self.w2_lora_a_stacked[0].shape[-2]
+
+ @classmethod
+ def can_replace_layer(
+ cls,
+ source_layer: nn.Module,
+ lora_config: LoRAConfig,
+ packed_modules_list: list,
+ model_config: PretrainedConfig | None,
+ ) -> bool:
+ """Returns True if the layer can be replaced by this LoRA layer."""
+
+ return type(source_layer) is FusedMoE and len(packed_modules_list) == 1
diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py
index adc5e861f57fb..c01984db4e64c 100644
--- a/vllm/lora/layers/logits_processor.py
+++ b/vllm/lora/layers/logits_processor.py
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import math
import torch
import torch.nn as nn
@@ -108,22 +107,13 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
(
max_loras,
1,
- # Pad for kernel compatibility
- math.ceil(
- self.base_layer.vocab_size / lora_config.lora_vocab_padding_size
- )
- * lora_config.lora_vocab_padding_size,
+ self.base_layer.vocab_size,
lora_config.max_lora_rank,
),
dtype=lora_config.lora_dtype,
device=self.device,
)
- self.embeddings_tensors = torch.full(
- (max_loras, lora_config.lora_extra_vocab_size, self.hidden_size),
- fill_value=float("-inf"),
- dtype=self.dtype,
- device=self.device,
- )
+
if self.sharded_to_full_mapping is not None:
self.sharded_to_full_mapping_gpu = torch.tensor(
self.sharded_to_full_mapping, device=self.device, dtype=torch.long
@@ -134,15 +124,15 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
def reset_lora(self, index: int):
self.lora_a_stacked[index] = 0
self.lora_b_stacked[index] = 0
- self.embeddings_tensors[index] = float("-inf")
def set_lora(
self,
index: int,
- lora_a: torch.Tensor,
- lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
):
+ assert isinstance(lora_a, torch.Tensor)
+ assert isinstance(lora_b, torch.Tensor)
self.reset_lora(index)
self.lora_a_stacked[index, 0, : lora_a.shape[0], : lora_a.shape[1]].copy_(
lora_a, non_blocking=True
@@ -150,12 +140,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
self.lora_b_stacked[index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_(
lora_b, non_blocking=True
)
- if embeddings_tensor is not None:
- self.embeddings_tensors[
- index,
- : embeddings_tensor.shape[0],
- : embeddings_tensor.shape[1],
- ] = embeddings_tensor
def _get_logits(
self,
@@ -193,39 +177,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
# token_id: [0, 1, 2, 3, 4, 5, -1, -1]
logits = logits[:, self.sharded_to_full_mapping_gpu]
- lora_logits = torch.empty(
- self.embeddings_tensors.shape[0] + 1,
- self.embeddings_tensors.shape[1],
- hidden_states.shape[0],
- dtype=self.embeddings_tensors.dtype,
- device=self.embeddings_tensors.device,
- )
- torch.matmul(self.embeddings_tensors, hidden_states.T, out=lora_logits[:-1])
-
- neg_inf, pos_inf = current_platform.get_infinity_values(lora_logits.dtype)
-
- lora_logits[-1] = neg_inf
- lora_logits = lora_logits.mT
- indices_padded = self.punica_wrapper.sampler_indices_padded
-
- if current_platform.is_tpu() or current_platform.is_xpu():
- indices_padded = indices_padded[: logits.size(0)]
-
- lora_logits = (
- lora_logits.reshape(
- lora_logits.shape[0] * lora_logits.shape[1],
- lora_logits.shape[2],
- )
- .index_select(0, indices_padded)
- .nan_to_num_(nan=neg_inf, posinf=pos_inf, neginf=neg_inf)
- )
-
- logits[
- :,
- self.base_layer.org_vocab_size : self.base_layer.org_vocab_size
- + lora_logits.shape[1],
- ] = lora_logits
-
lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits(
logits, hidden_states, self.lora_a_stacked, self.lora_b_stacked, 1.0
)
diff --git a/vllm/lora/layers/row_parallel_linear.py b/vllm/lora/layers/row_parallel_linear.py
index 2ef1bd98fc612..95517b1aee263 100644
--- a/vllm/lora/layers/row_parallel_linear.py
+++ b/vllm/lora/layers/row_parallel_linear.py
@@ -63,23 +63,18 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
input_parallel = splitted_input[self.tp_rank].contiguous()
# Matrix multiply.
- output_parallel = self.apply(input_parallel)
+ bias_ = (
+ None
+ if (self.tp_rank > 0 or self.base_layer.skip_bias_add)
+ else self.base_layer.bias
+ )
+ output_parallel = self.apply(input_parallel, bias_)
if self.base_layer.reduce_results and self.tp_size > 1:
- output_ = tensor_model_parallel_all_reduce(output_parallel)
+ output = tensor_model_parallel_all_reduce(output_parallel)
else:
- output_ = output_parallel
-
- if not self.base_layer.skip_bias_add:
- output = (
- output_ + self.base_layer.bias
- if self.base_layer.bias is not None
- else output_
- )
- output_bias = None
- else:
- output = output_
- output_bias = self.base_layer.bias
+ output = output_parallel
+ output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
if not self.base_layer.return_bias:
return output
@@ -120,7 +115,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
return lora_b
def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
- output = self.base_layer.quant_method.apply(self.base_layer, x)
+ output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
x = x.view(-1, x.shape[-1])
output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape
diff --git a/vllm/lora/layers/vocal_parallel_embedding.py b/vllm/lora/layers/vocal_parallel_embedding.py
index ca4ad8012e9c3..c87ca9e24dece 100644
--- a/vllm/lora/layers/vocal_parallel_embedding.py
+++ b/vllm/lora/layers/vocal_parallel_embedding.py
@@ -46,19 +46,10 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
self.embeddings_slice = None
self.embeddings_weights = None
- self.embeddings_tensors = torch.zeros(
- (
- max_loras,
- lora_config.lora_extra_vocab_size,
- self.base_layer.embedding_dim,
- ),
- dtype=self.base_layer.weight.dtype,
- device=self.base_layer.weight.device,
- )
self.lora_a_stacked = torch.zeros(
(
max_loras,
- self.base_layer.org_vocab_size + lora_config.lora_extra_vocab_size,
+ self.base_layer.org_vocab_size,
lora_config.max_lora_rank,
),
dtype=lora_config.lora_dtype,
@@ -82,54 +73,37 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
def reset_lora(self, index: int):
self.lora_a_stacked[index] = 0
self.lora_b_stacked[index] = 0
- self.embeddings_tensors[index] = 0
def set_lora(
self,
index: int,
- lora_a: torch.Tensor,
- lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None,
+ lora_a: torch.Tensor | list[torch.Tensor],
+ lora_b: torch.Tensor | list[torch.Tensor],
):
+ assert isinstance(lora_a, torch.Tensor)
+ assert isinstance(lora_b, torch.Tensor)
self.reset_lora(index)
# NOTE self.lora_a_stacked is row-major, and lora_a is col-major,
# so we need transpose here
+
self.lora_a_stacked[index, : lora_a.shape[1], : lora_a.shape[0]].copy_(
lora_a.T, non_blocking=True
)
self.lora_b_stacked[index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_(
lora_b, non_blocking=True
)
- if embeddings_tensor is not None:
- self.embeddings_tensors[
- index,
- : embeddings_tensor.shape[0],
- : embeddings_tensor.shape[1],
- ].copy_(embeddings_tensor, non_blocking=True)
- if self.embeddings_slice is not None:
- # TODO(yard1): Optimize this copy, we don't need to copy
- # everything, just the modified part
- embeddings = self.embeddings_tensors.view(
- self.embeddings_tensors.shape[0] * self.embeddings_tensors.shape[1],
- self.embeddings_tensors.shape[2],
- )[self.embeddings_slice[0] : self.embeddings_slice[1]]
- assert self.embeddings_weights is not None
- self.embeddings_weights[: embeddings.shape[0]].copy_(embeddings)
def forward(self, x: torch.Tensor) -> torch.Tensor:
- added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1, 1, 0)
-
# NB: Don't use torch.narrow here. torch.narrow triggers some
# Dynamic Shape specialization in torch.compile
num_tokens = x.shape[0]
indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
- indices_0 = self.punica_wrapper._embeddings_indices[0][:num_tokens]
full_lora_a_embeddings = F.embedding(
x + indices_1,
self.lora_a_stacked_2d,
)
- full_output = self.base_layer.forward(x + (indices_0 * added_tokens_mask))
+ full_output = self.base_layer.forward(x)
full_output_org = full_output
if full_output.ndim == 3:
diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py
index 7691481d5039e..f0d8e22194050 100644
--- a/vllm/lora/lora_weights.py
+++ b/vllm/lora/lora_weights.py
@@ -21,7 +21,6 @@ class LoRALayerWeights:
lora_alpha: int,
lora_a: torch.Tensor,
lora_b: torch.Tensor,
- embeddings_tensor: torch.Tensor | None = None,
scaling: float | None = None,
) -> None:
self.module_name = module_name
@@ -29,7 +28,6 @@ class LoRALayerWeights:
self.lora_alpha = lora_alpha
self.lora_a = lora_a
self.lora_b = lora_b
- self.embeddings_tensor = embeddings_tensor
if scaling is None:
self.scaling = self.lora_alpha / self.rank
@@ -56,18 +54,11 @@ class LoRALayerWeights:
def is_packed(self) -> bool:
return False
- @property
- def extra_vocab_size(self) -> int:
- return (
- self.embeddings_tensor.shape[0] if self.embeddings_tensor is not None else 0
- )
-
@classmethod
def from_config(
cls,
module_name: str,
peft_helper: PEFTHelper,
- embeddings_tensor: torch.Tensor | None = None,
) -> "LoRALayerWeights":
# lora_a and lora_b are set to None for config-based construction
return cls(
@@ -76,7 +67,6 @@ class LoRALayerWeights:
peft_helper.lora_alpha,
None,
None,
- embeddings_tensor,
peft_helper.vllm_lora_scaling_factor,
)
@@ -89,7 +79,6 @@ class LoRALayerWeights:
rank: int,
dtype: torch.dtype,
device: torch.types.Device,
- embeddings_tensor_dim: int | None = None,
) -> "LoRALayerWeights":
pin_memory = str(device) == "cpu" and is_pin_memory_available()
lora_a = torch.zeros(
@@ -99,24 +88,12 @@ class LoRALayerWeights:
[output_dim, rank], dtype=dtype, device=device, pin_memory=pin_memory
)
- embeddings_tensor = (
- torch.rand(
- 10,
- embeddings_tensor_dim,
- dtype=dtype,
- device=device,
- pin_memory=pin_memory,
- )
- if embeddings_tensor_dim
- else None
- )
return cls(
module_name,
rank=rank,
lora_alpha=1,
lora_a=lora_a,
lora_b=lora_b,
- embeddings_tensor=embeddings_tensor,
)
@@ -139,7 +116,6 @@ class PackedLoRALayerWeights(LoRALayerWeights):
lora_a=lora_a,
lora_b=lora_b,
scaling=scaling, # type: ignore
- embeddings_tensor=None,
)
self.lora_alphas = lora_alphas
if scaling is None:
diff --git a/vllm/lora/models.py b/vllm/lora/models.py
index 02c252f15bfab..636f062feb7b0 100644
--- a/vllm/lora/models.py
+++ b/vllm/lora/models.py
@@ -21,11 +21,14 @@ from vllm.lora.utils import (
from_layer,
from_layer_logits_processor,
get_supported_lora_modules,
+ is_base_embeddding_weights,
+ is_moe_model,
is_regex_target_modules,
parse_fine_tuned_lora_name,
process_packed_modules_mapping,
replace_submodule,
)
+from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from vllm.model_executor.models import SupportsLoRA, supports_multimodal
from vllm.model_executor.models.interfaces import is_pooling_model
@@ -93,14 +96,6 @@ class LoRAModel:
loras=self.loras.copy(),
)
- @property
- def extra_vocab_size(self) -> int:
- return (
- max(lora.extra_vocab_size for lora in self.loras.values())
- if self.loras
- else 0
- )
-
def get_lora(self, module_name: str) -> LoRALayerWeights | None:
"""Get LoRA for a given module by name"""
return self.loras.get(module_name, None)
@@ -117,7 +112,6 @@ class LoRAModel:
peft_helper: PEFTHelper,
device: str = "cuda",
dtype: torch.dtype | None = None,
- embeddings: dict[str, torch.Tensor] | None = None,
target_embedding_padding: int | None = None,
embedding_modules: dict[str, str] | None = None,
embedding_padding_modules: list[str] | None = None,
@@ -127,24 +121,14 @@ class LoRAModel:
pin_memory = str(device) == "cpu" and is_pin_memory_available()
loras: dict[str, LoRALayerWeights] = {}
for tensor_name, tensor in tensors.items():
+ if is_base_embeddding_weights(tensor_name):
+ continue
module_name, is_lora_a = parse_fine_tuned_lora_name(
tensor_name, weights_mapper
)
if module_name not in loras:
- lora_embeddings_tensor = None
- if embeddings:
- assert embedding_modules is not None
- embeddings_module = next(
- (k for k in embedding_modules if k in module_name), None
- )
- if embeddings_module:
- lora_embeddings_tensor = embeddings[
- embedding_modules[embeddings_module]
- ].to(device=device, dtype=dtype)
- if pin_memory:
- lora_embeddings_tensor = lora_embeddings_tensor.pin_memory()
loras[module_name] = LoRALayerWeights.from_config(
- module_name, peft_helper, lora_embeddings_tensor
+ module_name, peft_helper
)
if is_lora_a:
@@ -206,15 +190,17 @@ class LoRAModel:
lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors")
lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin")
lora_pt_file_path = os.path.join(lora_dir, "adapter_model.pt")
- new_embeddings_tensor_path = os.path.join(
- lora_dir, "new_embeddings.safetensors"
- )
- new_embeddings_bin_file_path = os.path.join(lora_dir, "new_embeddings.bin")
+ # new_embeddings_tensor_path = os.path.join(
+ # lora_dir, "new_embeddings.safetensors"
+ # )
+ # new_embeddings_bin_file_path = os.path.join(lora_dir, "new_embeddings.bin")
tensors: dict[str, torch.Tensor] = {}
unexpected_modules: list[list[str] | str] = []
def check_unexpected_modules(modules: dict):
for lora_module in modules.keys(): # noqa
+ if is_base_embeddding_weights(lora_module):
+ continue
module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
# Handle FSDP file format where experts.base_layer is the
# gate_up_proj and experts is the down_proj
@@ -300,21 +286,12 @@ class LoRAModel:
else:
raise ValueError(f"{lora_dir} doesn't contain tensors")
- embeddings = None
- if os.path.isfile(new_embeddings_tensor_path):
- embeddings = safetensors.torch.load_file(new_embeddings_tensor_path)
- elif os.path.isfile(new_embeddings_bin_file_path):
- embeddings = torch.load(
- new_embeddings_bin_file_path, map_location=device, weights_only=True
- )
-
return cls.from_lora_tensors(
lora_model_id=get_lora_id() if lora_model_id is None else lora_model_id,
tensors=tensors,
peft_helper=peft_helper,
device=device,
dtype=dtype,
- embeddings=embeddings,
target_embedding_padding=target_embedding_padding,
embedding_modules=embedding_modules,
embedding_padding_modules=embedding_padding_modules,
@@ -381,7 +358,11 @@ class LoRAModelManager:
self.modules: dict[str, BaseLayerWithLoRA] = {}
# Dict instead of a set for compatibility with LRUCache.
self._last_mapping: LoRAMapping | None = None
+ self._is_3d_moe_model = is_moe_model(self.model) and hasattr(
+ self.model, "is_3d_moe_weight"
+ )
self._create_lora_modules()
+
self.model.lora_manager = self
def __len__(self) -> int:
@@ -425,22 +406,36 @@ class LoRAModelManager:
self.lora_index_to_id[index] = lora_model.id
for module_name, module in self.modules.items():
module_lora = self._get_lora_layer_weights(lora_model, module_name)
- if module_lora:
- # Note (gnovack) - If MOE lora weights are not split into
- # num_experts chunks, we split them here
- if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(
- module_lora.lora_a
- ):
- # Handle FSDP file format where experts.base_layer is the
- # gate_up_proj and experts is the down_proj
- gate_up_proj_lora = self._get_lora_layer_weights(
- lora_model, module_name + ".base_layer"
- )
-
- assert gate_up_proj_lora is not None
- assert module_lora is not None
-
- down_proj_lora = module_lora
+ if not module_lora:
+ module.reset_lora(index)
+ continue
+ # Note (gnovack) - If MOE lora weights are not split into
+ # num_experts chunks, we split them here
+ if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(
+ module_lora.lora_a
+ ):
+ # Handle PEFT file format where experts.base_layer is the
+ # gate_up_proj and experts is the down_proj
+ gate_up_proj_lora = self._get_lora_layer_weights(
+ lora_model, module_name + ".base_layer"
+ )
+ down_proj_lora = module_lora
+ # FIXME Edge case where LoRA is not added to gate_up_proj
+ # or down_proj
+ assert gate_up_proj_lora is not None
+ assert down_proj_lora is not None
+ if self._is_3d_moe_model:
+ module_lora.lora_a = [
+ gate_up_proj_lora.lora_a,
+ down_proj_lora.lora_a,
+ ]
+ module_lora.lora_b = [
+ gate_up_proj_lora.lora_b,
+ down_proj_lora.lora_b,
+ ]
+ else:
+ # Some 3D MoE models haven't added the `is_3d_moe_weight`
+ # attribute yet, so fallback here
num_experts = module_lora.lora_a.shape[0] // module_lora.rank
gate_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
@@ -469,15 +464,12 @@ class LoRAModelManager:
module_lora.lora_a = lora_a
module_lora.lora_b = lora_b
+ module.set_lora(
+ index,
+ module_lora.lora_a,
+ module_lora.lora_b,
+ )
- module.set_lora(
- index,
- module_lora.lora_a,
- module_lora.lora_b,
- module_lora.embeddings_tensor,
- )
- else:
- module.reset_lora(index)
return True
def _deactivate_adapter(self, lora_id: int):
@@ -505,7 +497,6 @@ class LoRAModelManager:
self.lora_index_to_id,
self.lora_slots + 1,
self.vocab_size,
- self.lora_config.lora_extra_vocab_size,
)
def remove_all_adapters(self):
@@ -539,6 +530,13 @@ class LoRAModelManager:
continue
parts = module_name.split(".")[-1]
packed_moduled_lst = self.packed_modules_mapping.get(parts, [])
+ if isinstance(module, FusedMoE):
+ # packed_moduled_lst is used here to just determine whether to
+ # instantiate FusedMoE3DWithLoRA or FusedMoEWithLoRA, and the
+ # difference between these two LoRA layers is whether the
+ # LoRA weights of w1 and w3 have already been fused on disk.
+
+ packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"]
new_module = replace_submodule(
self.model,
module_name,
@@ -587,6 +585,7 @@ class LoRAModelManager:
self._register_packed_modules(module_name)
# All lora layers share the same punica_wrapper based on reference.
new_module.set_mapping(self.punica_wrapper)
+ pass
def register_module(self, module_name: str, module: "BaseLayerWithLoRA"):
assert isinstance(module, BaseLayerWithLoRA), (
@@ -616,7 +615,6 @@ class LoRAModelManager:
if parts[-1] in embedding_modules:
input_dim = (
module.base_layer.org_vocab_size
- + self.lora_config.lora_extra_vocab_size
if hasattr(module.base_layer, "org_vocab_size")
else module.base_layer.weight.shape[1]
)
@@ -625,11 +623,6 @@ class LoRAModelManager:
if hasattr(module.base_layer, "embedding_dim")
else module.base_layer.weight.shape[0]
)
- embeddings_tensor_dim = (
- module.base_layer.embedding_dim
- if hasattr(module.base_layer, "embedding_dim")
- else module.base_layer.weight.shape[1]
- )
lora = LoRALayerWeights.create_dummy_lora_weights(
module_name,
input_dim,
@@ -637,8 +630,31 @@ class LoRAModelManager:
rank,
module.lora_a_stacked[0].dtype,
"cpu",
- embeddings_tensor_dim=embeddings_tensor_dim,
)
+ model.loras[module_name] = lora
+ elif module.__class__.__name__ == "FusedMoE3DWithLoRA":
+ # Case for 3D moe model
+ # w2
+ lora = LoRALayerWeights.create_dummy_lora_weights(
+ module_name,
+ module.w2_input_size,
+ module.w2_output_size,
+ rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts
+ module.w2_lora_a_stacked[0].dtype,
+ "cpu",
+ )
+ model.loras[module_name] = lora
+ # w13
+ lora = LoRALayerWeights.create_dummy_lora_weights(
+ module_name,
+ module.w13_input_size,
+ module.w13_output_size,
+ rank
+ * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts
+ module.w13_lora_a_stacked[0].dtype,
+ "cpu",
+ )
+ model.loras[module_name + ".base_layer"] = lora
else:
lora = LoRALayerWeights.create_dummy_lora_weights(
module_name,
@@ -648,6 +664,7 @@ class LoRAModelManager:
module.lora_a_stacked[0].dtype,
"cpu",
)
+ model.loras[module_name] = lora
else:
parts = module_name.split(".")
replacements = self.packed_modules_mapping[parts[-1]]
@@ -663,7 +680,7 @@ class LoRAModelManager:
)
subloras.append(lora)
lora = PackedLoRALayerWeights.pack(subloras)
- model.loras[module_name] = lora
+ model.loras[module_name] = lora
return model
def _match_target_modules(self, module_name: str):
diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
index 893972144e99a..413ee8ecbbf96 100644
--- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+++ b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
@@ -3,6 +3,10 @@
import torch
+from vllm.distributed import (
+ tensor_model_parallel_all_gather,
+ tensor_model_parallel_all_reduce,
+)
from vllm.triton_utils import tl, triton
from vllm.utils.torch_utils import direct_register_custom_op
@@ -154,7 +158,7 @@ def _fused_moe_lora_kernel(
k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
# pre-fetch lora weight
b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
- # GDC wait waits for ALL programs in the the prior kernel to complete
+ # GDC wait waits for ALL programs in the prior kernel to complete
# before continuing.
if USE_GDC and not IS_PRIMARY:
tl.extra.cuda.gdc_wait()
@@ -311,6 +315,7 @@ def _fused_moe_lora_expand(
num_stages: int,
split_k: int,
mul_routed_weight: bool = False,
+ offset: int = 0,
) -> None:
b_ptr = _get_ptr(lora_b_stacked, device)
K = max_lora_rank
@@ -380,7 +385,7 @@ def _fused_moe_lora_expand(
**expand_config,
)
for i in range(num_slices):
- output[:, :, i * N : (i + 1) * N] += b_intermediate_cache1[i]
+ output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i]
@torch.inference_mode()
@@ -416,6 +421,8 @@ def _fused_moe_lora(
expand_num_stages: int,
expand_split_k: int,
mul_routed_weight: bool = False,
+ fully_sharded: bool = False,
+ offset: int = 0,
) -> None:
assert len(lora_a_stacked) == len(lora_b_stacked) > 0
assert (
@@ -430,7 +437,6 @@ def _fused_moe_lora(
== expert_ids.shape[0]
== num_tokens_post_padded.shape[0]
)
- assert len(lora_b_stacked) * lora_b_stacked[0].shape[-2] == output.shape[-1]
assert output.shape[0] == topk_weights.shape[0]
assert top_k_num == topk_weights.shape[1]
device = qcurr_hidden_states.device
@@ -480,6 +486,19 @@ def _fused_moe_lora(
mul_routed_weight,
)
+ if fully_sharded:
+ if max_lora_rank == w1_lora_b_stacked.shape[-1]:
+ a_intermediate_cache1 = tensor_model_parallel_all_reduce(
+ a_intermediate_cache1
+ )
+ else:
+ a_intermediate_cache1 = tensor_model_parallel_all_gather(
+ a_intermediate_cache1
+ )
+
+ # reset max_lora_rank to the full rank after allgather
+ max_lora_rank = a_intermediate_cache1.shape[-1]
+
_fused_moe_lora_expand(
output,
a_intermediate_cache1,
@@ -510,6 +529,7 @@ def _fused_moe_lora(
expand_num_stages,
expand_split_k,
mul_routed_weight,
+ offset,
)
diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py
index b6186e8561529..ce38751e4b6a7 100644
--- a/vllm/lora/punica_wrapper/punica_base.py
+++ b/vllm/lora/punica_wrapper/punica_base.py
@@ -31,7 +31,6 @@ class PunicaWrapperABC(ABC):
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
- extra_vocab_size: int,
**kwargs,
) -> None:
"""
@@ -172,8 +171,11 @@ class PunicaWrapperBase(PunicaWrapperABC):
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
- extra_vocab_size: int,
):
+ # NOTE We have remove lora extra vocab support for now. So we set
+ # extra_vocab_size alwayzs to 0, and extra_vocab_size will be removed.
+
+ extra_vocab_size = 0
(
base_indices,
sampler_indices,
@@ -285,12 +287,9 @@ class PunicaWrapperBase(PunicaWrapperABC):
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
- extra_vocab_size: int,
**kwargs,
):
- self._update_base_metadata(
- mapping, lora_index_to_id, max_loras, vocab_size, extra_vocab_size
- )
+ self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
if mapping.is_prefill:
# Update metadata required for prefill-related operators.
@@ -471,8 +470,8 @@ class PunicaWrapperBase(PunicaWrapperABC):
self,
y: torch.Tensor,
x: torch.Tensor,
- lora_a_stacked: list[torch.Tensor],
- lora_b_stacked: list[torch.Tensor],
+ lora_a_stacked: tuple[torch.Tensor, ...],
+ lora_b_stacked: tuple[torch.Tensor, ...],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
@@ -483,6 +482,8 @@ class PunicaWrapperBase(PunicaWrapperABC):
expand_config,
adapter_enabled: torch.Tensor,
mul_routed_weight=False,
+ fully_sharded: bool = False,
+ offset: int = 0,
):
"""
Performs a fused forward computation for LoRA of
diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py
index ede50a48af985..ef4b4ab7c3497 100644
--- a/vllm/lora/punica_wrapper/punica_gpu.py
+++ b/vllm/lora/punica_wrapper/punica_gpu.py
@@ -65,13 +65,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
- extra_vocab_size: int,
**kwargs,
):
self.is_prefill = mapping.is_prefill
- self._update_base_metadata(
- mapping, lora_index_to_id, max_loras, vocab_size, extra_vocab_size
- )
+ self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
# Prepare cuda kernel metadata tensors
self.token_mapping_meta.prepare_tensors(self.token_lora_indices)
@@ -363,8 +360,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
self,
y: torch.Tensor,
x: torch.Tensor,
- lora_a_stacked: list[torch.Tensor],
- lora_b_stacked: list[torch.Tensor],
+ lora_a_stacked: tuple[torch.Tensor, ...],
+ lora_b_stacked: tuple[torch.Tensor, ...],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
@@ -375,6 +372,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
expand_config,
adapter_enabled: torch.Tensor,
mul_routed_weight=False,
+ fully_sharded: bool = False,
+ offset: int = 0,
):
"""
Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
@@ -408,4 +407,6 @@ class PunicaWrapperGPU(PunicaWrapperBase):
expand_config.get("NUM_STAGES", 3),
expand_config.get("SPLIT_K", 1),
mul_routed_weight,
+ fully_sharded,
+ offset,
)
diff --git a/vllm/lora/punica_wrapper/punica_tpu.py b/vllm/lora/punica_wrapper/punica_tpu.py
index 090878dcd2546..0888772db54e7 100644
--- a/vllm/lora/punica_wrapper/punica_tpu.py
+++ b/vllm/lora/punica_wrapper/punica_tpu.py
@@ -292,7 +292,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
- extra_vocab_size: int,
):
# Make sure we don't accidentally collect outside operations
torch_xla.sync()
@@ -313,7 +312,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):
lora_index_to_id,
max_loras,
vocab_size,
- extra_vocab_size,
+ 0, # extra_vocab_size
"cpu",
)
self._token_lora_indices = self._pad_to_shape(
diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py
index b95087d0ff834..00c00782896cf 100644
--- a/vllm/lora/punica_wrapper/punica_xpu.py
+++ b/vllm/lora/punica_wrapper/punica_xpu.py
@@ -43,13 +43,10 @@ class PunicaWrapperXPU(PunicaWrapperBase):
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
- extra_vocab_size: int,
**kwargs,
):
self.is_prefill = mapping.is_prefill
- self._update_base_metadata(
- mapping, lora_index_to_id, max_loras, vocab_size, extra_vocab_size
- )
+ self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor:
return torch.narrow(self._token_lora_indices, 0, 0, x.size(0))
diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py
index 0f43ff06d8f2b..12524994d4968 100644
--- a/vllm/lora/utils.py
+++ b/vllm/lora/utils.py
@@ -23,6 +23,7 @@ from vllm.lora.layers import (
BaseLayerWithLoRA,
ColumnParallelLinearWithLoRA,
ColumnParallelLinearWithShardedLoRA,
+ FusedMoE3DWithLoRA,
FusedMoEWithLoRA,
LogitsProcessorWithLoRA,
MergedColumnParallelLinearWithLoRA,
@@ -62,6 +63,7 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {
MergedQKVParallelLinearWithShardedLoRA,
RowParallelLinearWithShardedLoRA,
FusedMoEWithLoRA,
+ FusedMoE3DWithLoRA,
}
@@ -166,6 +168,16 @@ def parse_fine_tuned_lora_name(
raise ValueError(f"{name} is unsupported LoRA weight")
+def is_base_embeddding_weights(name: str) -> bool:
+ # hardcoded subfixes for input & output embedding weights
+ input_embedding_subfix = ".embed_tokens.base_layer.weight"
+ output_embedding_subfix = ".lm_head.base_layer.weight"
+
+ return name.endswith(input_embedding_subfix) or name.endswith(
+ output_embedding_subfix
+ )
+
+
def is_regex_target_modules(
load_modules: str | list[str], expected_lora_modules: list[str]
) -> bool:
@@ -278,10 +290,12 @@ def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
# the expert indices are expanded based on the configured number
# of routed experts.
packed_modules_mapping = get_packed_modules_mapping(model)
-
- packed_modules_mapping["experts"] = [
- weight_name.rstrip(".") for _, weight_name, _, _ in moe_packed_mapping
- ]
+ if not hasattr(model, "is_3d_moe_weight"):
+ # 3D MoE LoRA does not need `packed_modules_mapping`
+ packed_modules_mapping["experts"] = [
+ weight_name.rstrip(".")
+ for _, weight_name, _, _ in moe_packed_mapping
+ ]
return packed_modules_mapping
else:
diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py
index b85151f2c7592..4cc201a6414f1 100644
--- a/vllm/lora/worker_manager.py
+++ b/vllm/lora/worker_manager.py
@@ -121,8 +121,7 @@ class WorkerLoRAManager:
lora_model_id=lora_request.lora_int_id,
device="cpu",
dtype=self.lora_config.lora_dtype,
- target_embedding_padding=self.vocab_size
- + self.lora_config.lora_extra_vocab_size,
+ target_embedding_padding=self.vocab_size,
embedding_modules=self.embedding_modules,
embedding_padding_modules=self.embedding_padding_modules,
tensorizer_config_dict=lora_request.tensorizer_config_dict,
@@ -143,12 +142,6 @@ class WorkerLoRAManager:
# For BadRequestError
raise e
- if lora.extra_vocab_size > self.lora_config.lora_extra_vocab_size:
- raise ValueError(
- f"LoRA added vocab size {lora.extra_vocab_size} "
- f"is greater than lora_extra_vocab_size "
- f"{self.lora_config.lora_extra_vocab_size}."
- )
return lora
def add_dummy_lora(self, lora_request: LoRARequest, rank: int) -> bool:
diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py
index 746a543ab827d..8b33727f05fbc 100644
--- a/vllm/model_executor/layers/batch_invariant.py
+++ b/vllm/model_executor/layers/batch_invariant.py
@@ -1,10 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import contextlib
import os
-from collections import namedtuple
from collections.abc import Callable
-from functools import cache
from typing import Any
import torch
@@ -725,10 +722,6 @@ _original_cublas_workspace_cfg = None
_original_cublaslt_workspace_size = None
-def is_batch_invariant_mode_enabled():
- return _batch_invariant_MODE
-
-
def enable_batch_invariant_mode():
global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
global _original_fp16_reduction_precision, _original_bf16_reduction_precision
@@ -791,83 +784,19 @@ def enable_batch_invariant_mode():
torch.backends.cuda.preferred_blas_library(backend="cublaslt")
-def disable_batch_invariant_mode():
- global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
- global _original_fp16_reduction_precision, _original_bf16_reduction_precision
- global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
- if not _batch_invariant_MODE:
- return
-
- if _batch_invariant_LIB is not None:
- _batch_invariant_LIB._destroy()
- if _original_torch_bmm is not None:
- torch.bmm = _original_torch_bmm
- _original_torch_bmm = None
-
- if _original_bf16_reduction_precision is not None:
- torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (
- _original_bf16_reduction_precision
- )
- _original_bf16_reduction_precision = None
- if _original_fp16_reduction_precision is not None:
- torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
- _original_fp16_reduction_precision
- )
- _original_fp16_reduction_precision = None
-
- torch.backends.cuda.preferred_blas_library(backend="default")
-
- if not is_torch_equal_or_newer("2.10.0.dev"):
- # Set cublas env vars to previous results. If previous results are None,
- # that means the env vars were not set, so we should remove them.
- if _original_cublas_workspace_cfg:
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = _original_cublas_workspace_cfg
- elif "CUBLAS_WORKSPACE_CONFIG" in os.environ:
- del os.environ["CUBLAS_WORKSPACE_CONFIG"]
-
- if _original_cublaslt_workspace_size:
- os.environ["CUBLASLT_WORKSPACE_SIZE"] = _original_cublaslt_workspace_size
- elif "CUBLASLT_WORKSPACE_SIZE" in os.environ:
- del os.environ["CUBLASLT_WORKSPACE_SIZE"]
-
- _original_cublas_workspace_cfg = None
- _original_cublaslt_workspace_size = None
-
- _batch_invariant_MODE = False
- _batch_invariant_LIB = None
-
-
-@contextlib.contextmanager
-def set_batch_invariant_mode(enabled: bool = True):
- global _batch_invariant_MODE, _batch_invariant_LIB
- old_data = (_batch_invariant_MODE, _batch_invariant_LIB)
- if enabled:
- enable_batch_invariant_mode()
- else:
- disable_batch_invariant_mode()
- yield
- if _batch_invariant_LIB is not None:
- _batch_invariant_LIB._destroy()
- _batch_invariant_MODE, _batch_invariant_LIB = old_data
-
-
-AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"])
-
-
-def get_batch_invariant_attention_block_size() -> AttentionBlockSize:
- return AttentionBlockSize(block_m=16, block_n=16)
-
-
-@cache
-def vllm_is_batch_invariant():
- env_key = "VLLM_BATCH_INVARIANT"
- is_overridden = False
- val = os.getenv(env_key, "0")
+def _read_vllm_batch_invariant() -> bool:
+ val = os.getenv("VLLM_BATCH_INVARIANT", "0")
try:
- is_overridden = int(val) != 0
+ return int(val) != 0
except ValueError:
- is_overridden = False
- return is_overridden
+ return False
+
+
+VLLM_BATCH_INVARIANT: bool = _read_vllm_batch_invariant()
+
+
+def vllm_is_batch_invariant() -> bool:
+ return VLLM_BATCH_INVARIANT
def override_envs_for_invariance():
@@ -876,11 +805,11 @@ def override_envs_for_invariance():
"FLASH_ATTN", # best supported backend
"FLASHINFER",
"FLASH_ATTN_MLA",
- "FLASHINFER_MLA",
- "TRITON_MLA",
# Not yet supported MLA backends
# "FLASHMLA",
# "FLEX_ATTENTION", # IMA issue even if we disable batch invariance
+ # "FLASHINFER_MLA", https://github.com/vllm-project/vllm/pull/28967
+ # "TRITON_MLA",
]
if curr_attn_backend not in supported_backends:
warning = (
@@ -923,5 +852,6 @@ def init_batch_invariance():
enable_batch_invariant_mode()
# Disable TF32 for batch invariance - it causes non-deterministic rounding
- torch.backends.cuda.matmul.allow_tf32 = False
- torch.backends.cudnn.allow_tf32 = False
+ torch.backends.cuda.matmul.fp32_precision = "ieee"
+ torch.backends.cudnn.conv.fp32_precision = "ieee"
+ torch.backends.cudnn.rnn.fp32_precision = "ieee"
diff --git a/vllm/model_executor/layers/conv.py b/vllm/model_executor/layers/conv.py
index e6f2d2990c241..8d51e5bd9920a 100644
--- a/vllm/model_executor/layers/conv.py
+++ b/vllm/model_executor/layers/conv.py
@@ -3,6 +3,7 @@
"""Conv Layer Class."""
import math
+from typing import Literal
import torch
import torch.nn as nn
@@ -23,11 +24,11 @@ class ConvLayerBase(CustomOp):
out_channels: int,
kernel_size: int | tuple[int, ...],
stride: int | tuple[int, ...] = 1,
- padding: int | tuple[int, ...] = 0,
+ padding: int | tuple[int, ...] | Literal["same", "valid"] = 0,
dilation: int | tuple[int, ...] = 1,
groups: int = 1,
bias: bool = True,
- padding_mode: str = "zeros",
+ padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
*,
params_dtype: torch.dtype | None = None,
) -> None:
@@ -36,6 +37,22 @@ class ConvLayerBase(CustomOp):
if params_dtype is None:
params_dtype = torch.get_default_dtype()
+ valid_padding_strings = {"same", "valid"}
+ if isinstance(padding, str) and padding not in valid_padding_strings:
+ raise ValueError(
+ f"Invalid padding string '{padding}'. "
+ f"Expected one of {valid_padding_strings}."
+ )
+
+ if padding == "same":
+ padding = (
+ kernel_size // 2
+ if isinstance(kernel_size, int)
+ else tuple(k // 2 for k in kernel_size)
+ )
+ elif padding == "valid":
+ padding = 0
+
kernel_size = (
(kernel_size,) * self.num_dim
if isinstance(kernel_size, int)
@@ -45,6 +62,9 @@ class ConvLayerBase(CustomOp):
padding = (padding,) * self.num_dim if isinstance(padding, int) else padding
dilation = (dilation,) * self.num_dim if isinstance(dilation, int) else dilation
+ if padding == "same" and any(s != 1 for s in stride):
+ raise ValueError("padding='same' is not supported for strided convolutions")
+
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py
index 2dd625054339c..86c50f39f0076 100644
--- a/vllm/model_executor/layers/fused_moe/all2all_utils.py
+++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py
@@ -67,6 +67,7 @@ def maybe_roundup_layer_hidden_size(
def maybe_make_prepare_finalize(
moe: FusedMoEConfig,
quant_config: FusedMoEQuantConfig | None,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> FusedMoEPrepareAndFinalize | None:
if not moe.moe_parallel_config.use_all2all_kernels:
return None
@@ -134,6 +135,13 @@ def maybe_make_prepare_finalize(
elif moe.use_deepep_ll_kernels:
assert quant_config is not None
+ global_to_physical = physical_to_global = local_expert_global_ids = None
+ if routing_tables is not None:
+ (
+ global_to_physical,
+ physical_to_global,
+ local_expert_global_ids,
+ ) = routing_tables
all_to_all_args = dict(
max_num_tokens_per_dp_rank=moe.max_num_tokens,
token_hidden_size=moe.hidden_dim,
@@ -155,6 +163,9 @@ def maybe_make_prepare_finalize(
max_tokens_per_rank=moe.max_num_tokens,
num_dispatchers=all2all_manager.world_size,
use_fp8_dispatch=use_fp8_dispatch,
+ global_to_physical=global_to_physical,
+ physical_to_global=physical_to_global,
+ local_expert_global_ids=local_expert_global_ids,
)
return prepare_finalize
diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py
index a7bd64b1c65e9..1826fafa8c4f5 100644
--- a/vllm/model_executor/layers/fused_moe/config.py
+++ b/vllm/model_executor/layers/fused_moe/config.py
@@ -8,7 +8,11 @@ import torch
import vllm.envs as envs
from vllm.config import ParallelConfig
-from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
+from vllm.distributed import (
+ get_dp_group,
+ get_pcp_group,
+ get_tensor_model_parallel_rank,
+)
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
OCP_MX_DTYPES,
@@ -24,10 +28,11 @@ logger = init_logger(__name__)
if has_triton_kernels():
try:
from triton_kernels.matmul_ogs import PrecisionConfig
- except ImportError:
+ except (ImportError, AttributeError) as e:
logger.error(
"Failed to import Triton kernels. Please make sure your triton "
- "version is compatible."
+ "version is compatible. Error: %s",
+ e,
)
@@ -684,9 +689,11 @@ FUSED_MOE_UNQUANTIZED_CONFIG: FusedMoEQuantConfig = FusedMoEQuantConfig.make()
@dataclass
class FusedMoEParallelConfig:
tp_size: int
+ pcp_size: int
dp_size: int
ep_size: int
tp_rank: int
+ pcp_rank: int
dp_rank: int
ep_rank: int
@@ -713,19 +720,22 @@ class FusedMoEParallelConfig:
return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
@staticmethod
- def flatten_tp_across_dp(
- tp_size: int, dp_size: int, dp_rank: int
+ def flatten_tp_across_dp_and_pcp(
+ tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
) -> tuple[int, int]:
tp_rank = 0 if tp_size == 1 else get_tensor_model_parallel_rank()
- # There are actually dp_size * tp_size devices. Update tp_size
- # and tp_rank so we shard across all devices.
- flatten_tp_size = dp_size * tp_size
- flatten_tp_rank = dp_rank * tp_size + tp_rank
+ # There are actually dp_size * pcp_size * tp_size devices.
+ # Update tp_size and tp_rank so we shard across all devices.
+ flatten_tp_size = dp_size * pcp_size * tp_size
+ flatten_tp_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
return flatten_tp_size, flatten_tp_rank
@staticmethod
def make(
- tp_size_: int, dp_size_: int, vllm_parallel_config: ParallelConfig
+ tp_size_: int,
+ pcp_size_: int,
+ dp_size_: int,
+ vllm_parallel_config: ParallelConfig,
) -> "FusedMoEParallelConfig":
"""
Determine MoE parallel configuration. Based on the input `tp_size_`,
@@ -734,19 +744,22 @@ class FusedMoEParallelConfig:
Args:
tp_size_ (int): `tp_size` passed into the FusedMoE constructor.
+ pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor.
dp_size_ (int): `dp_size` passed into the FusedMoE constructor.
vllm_parallel_config (ParallelConfig): vLLM's parallel config
object which contains the `enable_expert_parallel` flag.
Examples:
When there is no parallelism requested,
- i.e. `tp_size_` = `dp_size_` = 1, we simply return the sizes
+ i.e. `tp_size_` = `pcp_size_` = `dp_size_` = 1, we simply return the sizes
unaltered and the ranks set to 0.
- Expert Parallelism is considered only when either `dp_size_` or
+ Expert Parallelism is considered only when either `dp_size_`, `pcp_size_` or
`tp_size_` is non trivial.
- When TP = 2, DP = 1 and EP = False, the configuration on different
+ Note that PCP serves the same function as DP here.
+
+ When TP = 2, DP(PCP) = 1 and EP = False, the configuration on different
devices:
- device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} //
@@ -754,7 +767,7 @@ class FusedMoEParallelConfig:
- device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0}
- Comment : Tensors are sharded across 2 devices.
- When TP = 1, DP = 2 and EP = False, the configuration on different
+ When TP = 1, DP(PCP) = 2 and EP = False, the configuration on different
devices:
- device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0}
@@ -762,7 +775,7 @@ class FusedMoEParallelConfig:
- Comment: There are 2 engine instances and the tensors are sharded
across 2 decvices.
- When TP = 2, DP = 2 and EP = False, the configuration on different
+ When TP = 2, DP(PCP) = 2 and EP = False, the configuration on different
devices:
- device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0}
@@ -772,14 +785,14 @@ class FusedMoEParallelConfig:
- Comment: There are 2 engine instances and the tensors are sharded
across 4 devices.
- When, TP = 2, DP = 1 and EP = True, the configuration on different
+ When, TP = 2, DP(PCP) = 1 and EP = True, the configuration on different
devices:
- device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0}
- device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1}
- Comment: The experts are split between the 2 devices.
- When, TP = 1, DP = 2 and EP = True, the configuration on different
+ When, TP = 1, DP(PCP) = 2 and EP = True, the configuration on different
devices:
- device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0}
@@ -787,7 +800,7 @@ class FusedMoEParallelConfig:
- Comment: There are 2 engine instances and the experts are split
between the 2 devices.
- When TP = 2, DP = 2 and EP = True, the configuration on different
+ When TP = 2, DP(PCP) = 2 and EP = True, the configuration on different
devices:
- device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0}
@@ -798,18 +811,25 @@ class FusedMoEParallelConfig:
between the 4 devices.
"""
- use_ep = dp_size_ * tp_size_ > 1 and vllm_parallel_config.enable_expert_parallel
+ use_ep = (
+ dp_size_ * pcp_size_ * tp_size_ > 1
+ and vllm_parallel_config.enable_expert_parallel
+ )
dp_size = dp_size_
dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
- tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
- tp_size_, dp_size_, dp_rank
+ pcp_size = pcp_size_
+ pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
+ tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
+ tp_size_, dp_size_, dp_rank, pcp_size_, pcp_rank
)
if not use_ep:
return FusedMoEParallelConfig(
tp_size=tp_size,
tp_rank=tp_rank,
+ pcp_size=pcp_size,
+ pcp_rank=pcp_rank,
dp_size=dp_size,
dp_rank=dp_rank,
ep_size=1,
@@ -826,6 +846,8 @@ class FusedMoEParallelConfig:
return FusedMoEParallelConfig(
tp_size=1,
tp_rank=0,
+ pcp_size=pcp_size,
+ pcp_rank=pcp_rank,
dp_size=dp_size,
dp_rank=dp_rank,
ep_size=ep_size,
diff --git a/vllm/model_executor/layers/fused_moe/configs/E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/vllm/model_executor/layers/fused_moe/configs/E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json
new file mode 100644
index 0000000000000..54fe5374cb95d
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/configs/E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json
@@ -0,0 +1,147 @@
+{
+ "triton_version": "3.5.0",
+ "1": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 64,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "2": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "4": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "8": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "16": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "24": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "32": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "48": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "64": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "96": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "128": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "256": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "512": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 2
+ },
+ "1024": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 2
+ },
+ "1536": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "2048": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "3072": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 2
+ },
+ "4096": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ }
+}
diff --git a/vllm/model_executor/layers/fused_moe/configs/E=20,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8.json b/vllm/model_executor/layers/fused_moe/configs/E=20,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8.json
new file mode 100644
index 0000000000000..8b78f87e7f73b
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/configs/E=20,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8.json
@@ -0,0 +1,147 @@
+{
+ "triton_version": "3.5.0",
+ "1": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 2
+ },
+ "2": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 2
+ },
+ "8": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "16": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "24": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "32": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "48": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "64": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "96": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "128": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 256,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "256": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "512": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "1024": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 64,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "1536": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 64,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "2048": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 256,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "3072": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 256,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4096": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 256,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ }
+}
diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
index 6fcf408755f5d..532c16e899269 100644
--- a/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
+++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
@@ -1,11 +1,11 @@
{
"1": {
"BLOCK_SIZE_M": 16,
- "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 1,
+ "GROUP_SIZE_M": 16,
"num_warps": 4,
- "num_stages": 5
+ "num_stages": 4
},
"2": {
"BLOCK_SIZE_M": 16,
@@ -13,82 +13,82 @@
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
- "num_stages": 3
+ "num_stages": 4
},
"4": {
- "BLOCK_SIZE_M": 64,
- "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 4
},
"8": {
- "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "16": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "24": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "32": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "48": {
+ "BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3
},
- "16": {
- "BLOCK_SIZE_M": 16,
- "BLOCK_SIZE_N": 256,
- "BLOCK_SIZE_K": 64,
- "GROUP_SIZE_M": 32,
- "num_warps": 4,
- "num_stages": 3
- },
- "24": {
- "BLOCK_SIZE_M": 64,
- "BLOCK_SIZE_N": 128,
- "BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 16,
- "num_warps": 4,
- "num_stages": 3
- },
- "32": {
- "BLOCK_SIZE_M": 64,
- "BLOCK_SIZE_N": 128,
- "BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 1,
- "num_warps": 4,
- "num_stages": 3
- },
- "48": {
- "BLOCK_SIZE_M": 64,
- "BLOCK_SIZE_N": 128,
- "BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 32,
- "num_warps": 4,
- "num_stages": 3
- },
"64": {
- "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 32,
+ "GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3
},
"96": {
- "BLOCK_SIZE_M": 64,
- "BLOCK_SIZE_N": 128,
- "BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 32,
- "num_warps": 4,
- "num_stages": 3
- },
- "128": {
- "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
},
+ "128": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
"256": {
- "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
@@ -96,10 +96,10 @@
"num_stages": 3
},
"512": {
- "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 32,
+ "GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
},
@@ -109,7 +109,7 @@
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
- "num_stages": 3
+ "num_stages": 4
},
"1536": {
"BLOCK_SIZE_M": 64,
@@ -117,21 +117,21 @@
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
- "num_stages": 3
+ "num_stages": 4
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 64,
+ "GROUP_SIZE_M": 32,
"num_warps": 4,
- "num_stages": 3
+ "num_stages": 4
},
"3072": {
- "BLOCK_SIZE_M": 128,
- "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 16,
+ "GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 4
},
@@ -139,7 +139,7 @@
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
- "GROUP_SIZE_M": 64,
+ "GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
}
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 23ace3408562a..572307052b489 100644
--- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
@@ -6,7 +6,6 @@ import torch
from torch.nn import functional as F
from vllm import _custom_ops as ops
-from vllm import envs
def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
@@ -130,54 +129,6 @@ def select_experts(
)
-class IPEXFusedMOE:
- def __init__(self, layer: torch.nn.Module) -> None:
- import intel_extension_for_pytorch as ipex
-
- layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
- layer.w13_weight,
- layer.w2_weight,
- use_prepack=envs.VLLM_CPU_MOE_PREPACK,
- )
-
- def __call__(
- self,
- layer: torch.nn.Module,
- x: torch.Tensor,
- use_grouped_topk: bool,
- top_k: int,
- router_logits: torch.Tensor,
- renormalize: bool,
- topk_group: int | None = None,
- num_expert_group: int | None = None,
- global_num_experts: int = -1,
- expert_map: torch.Tensor | None = None,
- custom_routing_function: Callable | None = None,
- scoring_func: str = "softmax",
- routed_scaling_factor: float = 1.0,
- e_score_correction_bias: torch.Tensor | None = None,
- apply_router_weight_on_input: bool = False,
- activation: str = "silu",
- ) -> torch.Tensor:
- assert activation == "silu", f"{activation} is not supported."
- assert not apply_router_weight_on_input
- assert routed_scaling_factor == 1.0, (
- f"routed_scaling_factor {routed_scaling_factor} is not supported."
- )
- return layer.ipex_fusion(
- x,
- use_grouped_topk,
- top_k,
- router_logits,
- renormalize,
- topk_group,
- num_expert_group,
- custom_routing_function,
- scoring_func,
- e_score_correction_bias,
- )
-
-
class SGLFusedMOE:
def __init__(self, layer: torch.nn.Module) -> None:
pass
diff --git a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
index 06c9df317f7c7..fea9f49c04b89 100644
--- a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
@@ -6,6 +6,7 @@ import deep_ep
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm import envs
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -27,6 +28,8 @@ logger = init_logger(__name__)
DEEPEP_QUANT_BLOCK_SIZE = 128
DEEPEP_QUANT_BLOCK_SHAPE = [DEEPEP_QUANT_BLOCK_SIZE, DEEPEP_QUANT_BLOCK_SIZE]
+logger = init_logger(__name__)
+
def dequant_fp8(
expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor
@@ -85,6 +88,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
max_tokens_per_rank: int,
num_dispatchers: int,
use_fp8_dispatch: bool = False,
+ global_to_physical: torch.Tensor | None = None,
+ physical_to_global: torch.Tensor | None = None,
+ local_expert_global_ids: torch.Tensor | None = None,
):
super().__init__()
@@ -97,6 +103,17 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
self.handles: list[tuple | None] = [None, None]
self.num_dispatchers_ = num_dispatchers
+ topk_indices_dtype = self.topk_indices_dtype()
+
+ def _maybe_cast(tensor: torch.Tensor | None) -> torch.Tensor | None:
+ if tensor is None or topk_indices_dtype is None:
+ return tensor
+ return tensor.to(dtype=topk_indices_dtype)
+
+ self.global_to_physical = _maybe_cast(global_to_physical)
+ self.physical_to_global = _maybe_cast(physical_to_global)
+ self.local_expert_global_ids = _maybe_cast(local_expert_global_ids)
+
# We don't have enough information to determine if we should dispatch
# activation scales in a packed ue8m0 format during object construction
# time. This setting is handled by post_init_setup.
@@ -136,6 +153,16 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
def topk_indices_dtype(self) -> torch.dtype | None:
return torch.int64
+ def _map_global_to_physical_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
+ if self.global_to_physical is None:
+ return topk_ids
+ return self.global_to_physical[topk_ids]
+
+ def _map_local_to_global_ids(self, expert_topk_ids: torch.Tensor) -> torch.Tensor:
+ if self.local_expert_global_ids is None:
+ return expert_topk_ids
+ return self.local_expert_global_ids[expert_topk_ids]
+
def _do_quant(
self,
x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -163,16 +190,25 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
# TODO (varun): Optimization - Use a batched version of quant
x = x.view((-1, hidden_dim))
+ q_dtype = quant_config.quant_dtype
+
+ if envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm":
+ logger.info_once(
+ "Skip quantization when using FlashInfer CUTEDSL(masked_gemm) "
+ "for ModelOptNvFp4FusedMoE."
+ )
+ q_dtype = None
+
x, x_scales = moe_kernel_quantize_input(
x,
quant_config.a1_scale,
- quant_config.quant_dtype,
+ q_dtype,
quant_config.per_act_token_quant,
quant_config.block_shape,
)
x = x.view((num_experts, -1, hidden_dim))
- if quant_config.quant_dtype is not None:
+ if q_dtype is not None:
assert x_scales is not None
x_scales = normalize_batched_scales_shape(x_scales, num_experts)
@@ -226,9 +262,10 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
a1 = a1 * topk_weights.to(a1.dtype)
# Dispatch
+ dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)
expert_x, expert_num_tokens, handle, _, hook = self.buffer.low_latency_dispatch(
a1,
- topk_ids,
+ dispatch_topk_ids,
self.max_tokens_per_rank,
num_experts,
use_fp8=self.use_fp8_dispatch,
@@ -313,11 +350,12 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
# weights have already been applied.
combine_topk_weights = torch.ones_like(topk_weights)
+ combine_topk_ids = self._map_global_to_physical_ids(topk_ids)
# TODO (varun) : Enable zero copy mode
dbo_maybe_run_recv_hook()
_, _, recv_hook = self.buffer.low_latency_combine(
fused_expert_output,
- topk_ids,
+ combine_topk_ids,
combine_topk_weights,
handle,
async_finish=False,
diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py
new file mode 100644
index 0000000000000..2747ef04a3499
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py
@@ -0,0 +1,346 @@
+# 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.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+ TopKWeightAndReduceDelegate,
+)
+from vllm.utils.flashinfer import (
+ flashinfer_cutedsl_grouped_gemm_nt_masked,
+ has_flashinfer_cutedsl_grouped_gemm_nt_masked,
+ scaled_fp4_grouped_quantize,
+ silu_and_mul_scaled_nvfp4_experts_quantize,
+)
+
+logger = init_logger(__name__)
+
+
+def is_valid_flashinfer_cutedsl_fused_moe(
+ hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor
+) -> bool:
+ """
+ Check if the given problem size is supported by the FlashInfer CuteDSL MoE
+ kernel.
+ """
+ if not has_flashinfer_cutedsl_grouped_gemm_nt_masked():
+ logger.debug_once(
+ "FlashInferCuteDSLExperts disabled: "
+ "flashinfer_cutedsl_fused_moe not available."
+ )
+ return False
+ # Data type checks
+ if (
+ w1.dtype != torch.uint8
+ or w2.dtype != torch.uint8
+ or hidden_states.dtype not in [torch.float32, torch.float16, torch.bfloat16]
+ ):
+ logger.debug_once(
+ "FlashInferCuteDSLExperts disabled: w1/w2 must be torch.uint8 "
+ f"(got w1={w1.dtype}, w2={w2.dtype}), hidden_states must be "
+ f"float32, float16, or bfloat16 (got {hidden_states.dtype})."
+ )
+ return False
+ return True
+
+
+class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):
+ def __init__(
+ self,
+ out_dtype: torch.dtype,
+ quant_config: FusedMoEQuantConfig,
+ ):
+ super().__init__(quant_config)
+ assert quant_config.quant_dtype == "nvfp4", (
+ "Only nvfp4 quantization are currently supported."
+ )
+ self.out_dtype = out_dtype
+
+ @property
+ def activation_formats(
+ self,
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
+ return (
+ mk.FusedMoEActivationFormat.BatchedExperts,
+ mk.FusedMoEActivationFormat.BatchedExperts,
+ )
+
+ def supports_expert_map(self) -> bool:
+ return False
+
+ def supports_chunking(self) -> bool:
+ # This refers to TP chunking; DP chunking is handled separately.
+ # TODO(shuw@nvidia.com): Set to False to be consistent with
+ # batched_deep_gemm_moe
+ return False
+
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+ # Let PrepareAndFinalize::finalize() decide the impl.
+ return TopKWeightAndReduceDelegate()
+
+ 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,
+ ) -> 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.
+ """
+ output_shape = (local_num_experts, M, K)
+ workspace2 = (local_num_experts, M, N)
+ workspace1 = output_shape
+ 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: str,
+ global_num_experts: int,
+ expert_map: torch.Tensor | None,
+ a1q_scale: torch.Tensor | None,
+ a2_scale: torch.Tensor | None, # Not used
+ workspace13: torch.Tensor | None,
+ workspace2: torch.Tensor | None,
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
+ apply_router_weight_on_input: bool | None,
+ ):
+ assert self.quant_dtype == "nvfp4", (
+ "Only nvfp4 quantization are currently supported."
+ )
+ # Ensure w1_scale and w2_scale are not None before calling view
+ assert self.w1_scale is not None and self.w2_scale is not None, (
+ "w1_scale and w2_scale must not be None for FlashInferExperts"
+ )
+ assert expert_tokens_meta is not None
+ expert_num_tokens = expert_tokens_meta.expert_num_tokens
+ assert hidden_states.ndim == 3
+ assert self.w1_scale.ndim == 3
+ assert self.w2_scale.ndim == 3
+ flashinfer_cutedsl_moe_masked(
+ hidden_states=hidden_states,
+ input_global_scale=self.a1_gscale,
+ w1=w1,
+ w1_blockscale=self.w1_scale,
+ w1_alpha=self.g1_alphas,
+ w2=w2,
+ a2_global_scale=self.a2_gscale,
+ w2_blockscale=self.w2_scale,
+ w2_alpha=self.g2_alphas,
+ masked_m=expert_num_tokens,
+ workspace=workspace2,
+ out=output,
+ )
+
+
+def get_cute_dtype(input: torch.Tensor) -> str:
+ if input.dtype == torch.bfloat16:
+ return "bfloat16"
+ elif input.dtype == torch.float16:
+ return "float16"
+ elif input.dtype == torch.float32:
+ return "float32"
+ else:
+ raise ValueError(f"Unsupported cute dtype {input.dtype}")
+
+
+def flashinfer_cutedsl_moe_masked(
+ hidden_states: torch.Tensor,
+ input_global_scale: torch.Tensor,
+ w1: torch.Tensor,
+ w1_blockscale: torch.Tensor,
+ w1_alpha,
+ w2: torch.Tensor,
+ a2_global_scale: torch.Tensor,
+ w2_blockscale: torch.Tensor,
+ w2_alpha,
+ masked_m: torch.Tensor,
+ workspace: torch.Tensor,
+ out: torch.Tensor,
+):
+ """
+ Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
+ kernels.
+
+ Args:
+ hidden_states (torch.Tensor): [num_experts, m, k], bf16
+ input_global_scale (torch.Tensor): (l,)
+ w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8
+ w1_blockscale (torch.Tensor): blockscale factors, e4m3,
+ w1_alpha (torch.Tensor): (l,)
+ w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8
+ a2_global_scale (torch.Tensor): (l,)
+ w2_blockscale (torch.Tensor): blockscale factors, e4m3,
+ w2_alpha (torch.Tensor): (l,)
+ masked_m (torch.Tensor): Masked dimension indices
+ workspace (torch.Tensor): For gateup_output
+
+ Notes:
+ - Assumes max(masked_m) <= m.
+ """
+
+ # === Assertions on dtypes ===
+ assert input_global_scale.dtype == torch.float32, (
+ f"input_global_scale must be float32, got {input_global_scale.dtype}"
+ )
+ assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
+ assert w1_blockscale.dtype == torch.float8_e4m3fn, (
+ f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}"
+ )
+ assert w1_alpha.dtype == torch.float32, (
+ f"w1_alpha must be float32, got {w1_alpha.dtype}"
+ )
+ assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}"
+ assert a2_global_scale.dtype == torch.float32, (
+ f"a2_global_scale must be float32, got {a2_global_scale.dtype}"
+ )
+ assert w2_blockscale.dtype == torch.float8_e4m3fn, (
+ f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}"
+ )
+ assert w2_alpha.dtype == torch.float32, (
+ f"w2_alpha must be float32, got {w2_alpha.dtype}"
+ )
+
+ # === Assertions on shapes ===
+ n = w2.shape[-1] * 2 # intermediate dimension
+ num_experts, m, k = hidden_states.shape
+
+ assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
+ assert w1.shape[-1] * 2 == k, (
+ f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}"
+ )
+ assert w2.shape[-2:] == (
+ k,
+ n // 2,
+ ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}"
+
+ assert input_global_scale.shape == (num_experts,), (
+ f"input_global_scale must be (l,), got {input_global_scale.shape}"
+ )
+ assert w1_alpha.shape == (num_experts,), (
+ f"w1_alpha must be (l,), got {w1_alpha.shape}"
+ )
+ assert a2_global_scale.shape == (num_experts,), (
+ f"a2_global_scale must be (l,), got {a2_global_scale.shape}"
+ )
+ assert w2_alpha.shape == (num_experts,), (
+ f"w2_alpha must be (l,), got {w2_alpha.shape}"
+ )
+
+ aq, aq_sf = scaled_fp4_grouped_quantize(
+ hidden_states,
+ masked_m,
+ input_global_scale,
+ )
+
+ workspace = workspace.permute(1, 2, 0) # requirement of kernel
+ sf_vec_size = 16
+ assert aq_sf.dtype == torch.float8_e4m3fn
+ assert aq.dtype == torch.uint8
+ ab_dtype = "float4_e2m1fn"
+ sf_dtype = "float8_e4m3fn"
+
+ c_dtype = get_cute_dtype(hidden_states)
+
+ # Gemm1
+ flashinfer_cutedsl_grouped_gemm_nt_masked(
+ (aq, aq_sf),
+ (w1.permute(1, 2, 0), w1_blockscale),
+ workspace,
+ masked_m,
+ ab_dtype=ab_dtype,
+ sf_dtype=sf_dtype,
+ c_dtype=c_dtype,
+ sf_vec_size=sf_vec_size,
+ alpha=w1_alpha.view(1, 1, num_experts),
+ alpha_dtype=get_cute_dtype(w1_alpha),
+ ) # in logical [m, n, l]
+
+ # SILU and quantization
+ diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
+ workspace.permute(2, 0, 1),
+ masked_m,
+ a2_global_scale,
+ )
+
+ # Gemm2
+ out = out.permute(1, 2, 0) # requirement of kernel
+ flashinfer_cutedsl_grouped_gemm_nt_masked(
+ (diq, diq_sf),
+ (w2.permute(1, 2, 0), w2_blockscale),
+ out,
+ masked_m,
+ ab_dtype=ab_dtype,
+ sf_dtype=sf_dtype,
+ c_dtype=c_dtype,
+ sf_vec_size=sf_vec_size,
+ alpha=w2_alpha.view(1, 1, num_experts),
+ alpha_dtype=get_cute_dtype(w2_alpha),
+ ) # in logical [m, k, l]
+ out = out.permute(2, 0, 1)
+
+
+def flashinfer_cutedsl_moe_fp4(
+ hidden_states: torch.Tensor,
+ w1: torch.Tensor,
+ w2: torch.Tensor,
+ topk_weights: torch.Tensor,
+ topk_ids: torch.Tensor,
+ quant_config: FusedMoEQuantConfig,
+ inplace: bool = False,
+ activation: str = "silu",
+ global_num_experts: int = -1,
+ expert_map: torch.Tensor | None = None,
+ apply_router_weight_on_input: bool = False,
+) -> torch.Tensor:
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
+ create_flashinfer_prepare_finalize,
+ )
+
+ fused_experts = mk.FusedMoEModularKernel(
+ create_flashinfer_prepare_finalize(use_dp=False), # could be swapped later
+ FlashInferCuteDSLExperts(
+ out_dtype=hidden_states.dtype,
+ quant_config=quant_config,
+ ),
+ )
+
+ return fused_experts(
+ hidden_states=hidden_states,
+ w1=w1,
+ w2=w2,
+ topk_weights=topk_weights,
+ topk_ids=topk_ids,
+ inplace=inplace,
+ activation=activation,
+ global_num_experts=global_num_experts,
+ expert_map=expert_map,
+ apply_router_weight_on_input=apply_router_weight_on_input,
+ )
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index 2e042d85fcfcf..df208eae2e71c 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -872,8 +872,10 @@ def get_moe_configs(
for config_file_path in config_file_paths:
if os.path.exists(config_file_path):
with open(config_file_path) as f:
- logger.info(
- "Using configuration from %s for MoE layer.", config_file_path
+ logger.info_once(
+ "Using configuration from %s for MoE layer.",
+ config_file_path,
+ scope="global",
)
# If a configuration has been found, return it
tuned_config = json.load(f)
@@ -1246,7 +1248,6 @@ def eplb_map_to_physical_and_record(
expert_load_view: torch.Tensor,
logical_to_physical_map: torch.Tensor,
logical_replica_count: torch.Tensor,
- indices_type: torch.dtype | None = None,
) -> torch.Tensor:
"""
Map the logical expert ids to physical expert ids
@@ -1260,7 +1261,6 @@ def eplb_map_to_physical_and_record(
expert_load_view: The expert load view.
logical_to_physical_map: The logical to physical map.
logical_replica_count: The logical replica count.
- indices_type: The indices type.
Returns:
The physical expert ids.
@@ -1310,9 +1310,6 @@ def eplb_map_to_physical_and_record(
index=topk_ids_flatten.long(),
src=torch.ones_like(topk_ids_flatten).to(expert_load_view),
)
-
- if indices_type is not None:
- topk_ids = topk_ids.to(dtype=indices_type)
return topk_ids
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
index 87f8c8d75a9b5..073e90a4e6808 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
@@ -50,10 +50,15 @@ class FusedMoEMethodBase(QuantizeMethodBase):
"""
return False
- def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
+ def maybe_make_prepare_finalize(
+ self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+ ) -> FusedMoEPrepareAndFinalize | None:
from .all2all_utils import maybe_make_prepare_finalize
- return maybe_make_prepare_finalize(self.moe, self.moe_quant_config)
+ return maybe_make_prepare_finalize(
+ self.moe, self.moe_quant_config, routing_tables
+ )
def select_gemm_impl(
self,
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
index 43974ba917e42..c6dc95acdb636 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
@@ -50,6 +50,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
prepare_finalize,
old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
shared_experts,
+ getattr(moe_layer, "shared_experts_stream", None),
),
)
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index 023132acfed3f..6619b64b2bbc0 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -5,7 +5,7 @@ from collections.abc import Callable, Iterable
from contextlib import nullcontext
from enum import Enum
from functools import partial
-from typing import Literal, get_args, overload
+from typing import Literal, cast, get_args, overload
import torch
import torch.nn.functional as F
@@ -18,6 +18,7 @@ from vllm.config.parallel import ExpertPlacementStrategy
from vllm.distributed import (
get_dp_group,
get_ep_group,
+ get_pcp_group,
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_reduce,
)
@@ -67,7 +68,6 @@ else:
expert_load_view: torch.Tensor,
logical_to_physical_map: torch.Tensor,
logical_replica_count: torch.Tensor,
- indices_type: torch.dtype | None,
) -> torch.Tensor:
# CPU fallback: no EPLB so just return as is
return topk_ids
@@ -192,6 +192,42 @@ def determine_expert_map(
return (local_num_experts, expert_map, expert_mask)
+def determine_expert_placement_strategy(
+ expert_placement_strategy: ExpertPlacementStrategy,
+ moe_parallel_config: FusedMoEParallelConfig,
+ num_expert_group: int | None,
+ num_redundant_experts: int,
+ enable_eplb: bool,
+) -> ExpertPlacementStrategy:
+ if expert_placement_strategy == "round_robin":
+ round_robin_supported = (
+ (num_expert_group is not None and num_expert_group > 1)
+ and num_redundant_experts == 0
+ and not enable_eplb
+ )
+
+ if not round_robin_supported:
+ logger.warning(
+ "Round-robin expert placement is only supported for "
+ "models with multiple expert groups and no redundant "
+ "experts. Falling back to linear expert placement."
+ )
+ return "linear"
+ if (
+ moe_parallel_config.use_all2all_kernels
+ and not moe_parallel_config.use_deepep_ll_kernels
+ ):
+ logger.warning(
+ "Round-robin expert placement currently only supports "
+ "the DeepEP low-latency backend, but '%s' was configured. "
+ "Falling back to linear expert placement.",
+ moe_parallel_config.all2all_backend,
+ )
+ return "linear"
+
+ return expert_placement_strategy
+
+
def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
"""
Compresses the expert map by removing any -1 entries.
@@ -307,6 +343,7 @@ class FusedMoE(CustomOp):
tp_size: int | None = None,
ep_size: int | None = None,
dp_size: int | None = None,
+ pcp_size: int | None = None,
prefix: str = "",
custom_routing_function: Callable | None = None,
scoring_func: str = "softmax",
@@ -335,8 +372,8 @@ class FusedMoE(CustomOp):
logger.info_once("Disabling MoE shared_experts cuda stream")
self.shared_experts_stream = None
else:
- # TODO(rob): enable shared expert overlap with non-cuda.
- # aux_stream() returns None on non-cuda platforms.
+ # TODO(rob): enable shared expert overlap with non-cuda-alike.
+ # aux_stream() returns None on non-cuda-alike platforms.
self.shared_experts_stream = aux_stream()
if self.shared_experts_stream is not None:
logger.info_once("Enabled separate cuda stream for MoE shared_experts")
@@ -362,12 +399,14 @@ class FusedMoE(CustomOp):
tp_size if tp_size is not None else get_tensor_model_parallel_world_size()
)
dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size
+ pcp_size_ = pcp_size if pcp_size is not None else get_pcp_group().world_size
self.is_sequence_parallel = is_sequence_parallel
self.sp_size = tp_size_ if is_sequence_parallel else 1
self.moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
tp_size_=tp_size_,
+ pcp_size_=pcp_size_,
dp_size_=dp_size_,
vllm_parallel_config=vllm_config.parallel_config,
)
@@ -400,6 +439,9 @@ class FusedMoE(CustomOp):
self.expert_load_view: torch.Tensor | None = None
self.logical_to_physical_map: torch.Tensor | None = None
self.logical_replica_count: torch.Tensor | None = None
+ self.expert_placement_strategy: ExpertPlacementStrategy = (
+ vllm_config.parallel_config.expert_placement_strategy
+ )
# ROCm aiter shared experts fusion
self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
@@ -433,38 +475,27 @@ class FusedMoE(CustomOp):
"Redundant experts are only supported with EPLB."
)
- expert_placement_strategy = (
- vllm_config.parallel_config.expert_placement_strategy
+ self.expert_placement_strategy = determine_expert_placement_strategy(
+ expert_placement_strategy=self.expert_placement_strategy,
+ moe_parallel_config=self.moe_parallel_config,
+ num_expert_group=num_expert_group,
+ num_redundant_experts=num_redundant_experts,
+ enable_eplb=self.enable_eplb,
)
- if expert_placement_strategy == "round_robin":
- # TODO(Bruce): will support round robin expert placement with
- # EPLB enabled in the future.
- round_robin_supported = (
- (num_expert_group is not None and num_expert_group > 1)
- and num_redundant_experts == 0
- and not self.enable_eplb
- )
-
- if not round_robin_supported:
- logger.warning(
- "Round-robin expert placement is only supported for "
- "models with multiple expert groups and no redundant "
- "experts. Falling back to linear expert placement."
- )
- expert_placement_strategy = "linear"
self.expert_map: torch.Tensor | None
local_num_experts, expert_map, expert_mask = determine_expert_map(
ep_size=self.ep_size,
ep_rank=self.ep_rank,
global_num_experts=self.global_num_experts,
- expert_placement_strategy=expert_placement_strategy,
+ expert_placement_strategy=self.expert_placement_strategy,
num_fused_shared_experts=self.num_fused_shared_experts,
return_expert_mask=self.rocm_aiter_fmoe_enabled,
)
self.local_num_experts = local_num_experts
self.register_buffer("expert_map", expert_map)
self.register_buffer("expert_mask", expert_mask)
+ self._maybe_init_expert_routing_tables()
logger.info_once(
"[EP Rank %s/%s] Expert parallelism is enabled. Expert "
"placement strategy: %s. Local/global"
@@ -472,7 +503,7 @@ class FusedMoE(CustomOp):
" %s.",
self.ep_rank,
self.ep_size,
- expert_placement_strategy,
+ self.expert_placement_strategy,
self.local_num_experts,
self.global_num_experts,
get_compressed_expert_map(self.expert_map),
@@ -542,6 +573,9 @@ class FusedMoE(CustomOp):
is_act_and_mul=is_act_and_mul,
is_lora_enabled=vllm_config.lora_config is not None,
)
+ self.moe_config_use_flashinfer_cutlass_kernels = (
+ self.moe_config.use_flashinfer_cutlass_kernels
+ )
self.quant_config = quant_config
@@ -621,7 +655,12 @@ class FusedMoE(CustomOp):
# should be safe to swap out the quant_method.
def maybe_init_modular_kernel(self) -> None:
self.ensure_moe_quant_config_init()
- prepare_finalize = self.quant_method.maybe_make_prepare_finalize()
+ # routing_tables only needed for round-robin expert placement with
+ # DeepEP all2all backend.
+ routing_tables = self._maybe_init_expert_routing_tables()
+ prepare_finalize = self.quant_method.maybe_make_prepare_finalize(
+ routing_tables=routing_tables
+ )
if prepare_finalize is not None:
logger.debug(
"%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)
@@ -646,6 +685,10 @@ class FusedMoE(CustomOp):
def dp_size(self):
return self.moe_parallel_config.dp_size
+ @property
+ def pcp_size(self):
+ return self.moe_parallel_config.pcp_size
+
@property
def ep_size(self):
return self.moe_parallel_config.ep_size
@@ -658,6 +701,10 @@ class FusedMoE(CustomOp):
def dp_rank(self):
return self.moe_parallel_config.dp_rank
+ @property
+ def pcp_rank(self):
+ return self.moe_parallel_config.pcp_rank
+
@property
def ep_rank(self):
return self.moe_parallel_config.ep_rank
@@ -683,7 +730,7 @@ class FusedMoE(CustomOp):
return (
self.moe_quant_config is not None
and self.moe_quant_config.quant_dtype == "nvfp4"
- and self.moe_config.use_flashinfer_cutlass_kernels
+ and self.moe_config_use_flashinfer_cutlass_kernels
)
@property
@@ -703,6 +750,84 @@ class FusedMoE(CustomOp):
# By default, router/gate is called before FusedMoE forward pass
return False
+ def _maybe_init_expert_routing_tables(
+ self,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+ # Currently routing_tables only needed for round-robin expert placement
+ # with DeepEP-ll all2all backend.
+ if (
+ self.expert_placement_strategy != "round_robin"
+ or not self.use_deepep_ll_kernels
+ ):
+ return None
+
+ if hasattr(self, "expert_global_to_physical"):
+ return cast(
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+ (
+ self.expert_global_to_physical,
+ self.expert_physical_to_global,
+ self.expert_local_to_global,
+ ),
+ )
+
+ if self.expert_map is None:
+ return None
+
+ routing_tables = self.ensure_round_robin_expert_routing_tables(
+ global_num_experts=self.global_num_experts,
+ ep_size=self.ep_size,
+ ep_rank=self.ep_rank,
+ local_num_experts=self.local_num_experts,
+ device=self.expert_map.device,
+ )
+
+ global_to_physical, physical_to_global, local_global = routing_tables
+ self.register_buffer("expert_global_to_physical", global_to_physical)
+ self.register_buffer("expert_physical_to_global", physical_to_global)
+ self.register_buffer("expert_local_to_global", local_global)
+
+ return routing_tables
+
+ @staticmethod
+ def ensure_round_robin_expert_routing_tables(
+ global_num_experts: int,
+ ep_size: int,
+ ep_rank: int,
+ local_num_experts: int,
+ device: torch.device | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ device_kwargs = {"device": device} if device is not None else {}
+ global_indices = torch.arange(
+ global_num_experts, dtype=torch.long, **device_kwargs
+ )
+ owner = torch.remainder(global_indices, ep_size)
+ local_index = torch.div(global_indices, ep_size, rounding_mode="floor")
+ base = global_num_experts // ep_size
+ remainder = global_num_experts % ep_size
+ physical_offset = owner * base
+ if remainder > 0:
+ remainder_tensor = torch.tensor(
+ remainder, dtype=torch.long, **device_kwargs
+ )
+ physical_offset = physical_offset + torch.minimum(owner, remainder_tensor)
+
+ global_to_physical = physical_offset + local_index
+ physical_to_global = torch.empty_like(global_to_physical)
+ physical_to_global[global_to_physical] = global_indices
+
+ local_global = torch.arange(
+ ep_rank,
+ global_num_experts,
+ ep_size,
+ dtype=torch.long,
+ **device_kwargs,
+ )
+ if local_global.numel() != local_num_experts:
+ local_global = local_global[:local_num_experts]
+
+ return (global_to_physical, physical_to_global, local_global)
+
def update_expert_map(self):
# ep_size and ep_rank should already be updated
assert self.expert_map is not None
@@ -711,18 +836,59 @@ class FusedMoE(CustomOp):
ep_size=self.ep_size,
ep_rank=self.ep_rank,
global_num_experts=self.global_num_experts,
+ expert_placement_strategy=self.expert_placement_strategy,
num_fused_shared_experts=self.num_fused_shared_experts,
return_expert_mask=self.rocm_aiter_fmoe_enabled,
)
self.local_num_experts = local_num_experts
self.register_buffer("expert_map", expert_map)
self.register_buffer("expert_mask", expert_mask)
+ self._maybe_init_expert_routing_tables()
if self.aiter_fmoe_shared_expert_enabled:
self._init_aiter_shared_experts_topK_buffer(
vllm_config=get_current_vllm_config(),
dp_size=get_dp_group().world_size,
)
+ def _maybe_setup_shared_experts_stream(
+ self,
+ hidden_states: torch.Tensor,
+ has_separate_shared_experts: bool,
+ use_chunked_impl: bool,
+ ) -> tuple[bool, torch.Tensor | None]:
+ use_shared_experts_stream = (
+ has_separate_shared_experts
+ and not use_chunked_impl
+ and self.shared_experts_stream is not None
+ and (
+ hidden_states.shape[0]
+ <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
+ )
+ )
+
+ hidden_states_clone: torch.Tensor | None = None
+ if use_shared_experts_stream:
+ assert self.shared_experts_stream is not None
+
+ # Clone BEFORE switching streams to avoid race condition
+ # where routed_expert kernel may mutate hidden_states.
+ hidden_states_clone = hidden_states.clone()
+
+ # Record that the clone will be used by shared_experts_stream
+ # to avoid gc issue from deallocation of hidden_states_clone
+ # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
+ # NOTE: We dont need shared_output.record_stream(current_stream())
+ # because we synch the streams before using shared_output.
+ hidden_states_clone.record_stream(self.shared_experts_stream)
+
+ # Mark sync start point for the separate shared experts
+ # stream here since we want to run in parallel with the
+ # router/gate (next op below)
+ assert self.shared_experts_stream is not None
+ self.shared_experts_stream.wait_stream(current_stream())
+
+ return use_shared_experts_stream, hidden_states_clone
+
def _load_per_tensor_weight_scale(
self,
shard_id: str,
@@ -1225,7 +1391,48 @@ class FusedMoE(CustomOp):
yield param_name
def get_expert_weights(self) -> Iterable[torch.Tensor]:
+ def _maybe_make_contiguous(
+ name: str, p: torch.nn.Parameter
+ ) -> torch.nn.Parameter:
+ """
+ In some cases, the last 2 dimensions (the non-expert dimensions)
+ of the weight scale tensor are transposed. This function
+ transforms the tensor (view update) so the tensor is contiguous().
+ Example: A non-contiguous scale tensor,
+ `x` of shape (E, 32, 16) and stride (512, 1, 32) is transformed to
+ `x_` of shape (E, 16, 32) and stride (512, 32, 1).
+ Note that we specifically use torch.transpose() so `x_` refers
+ to the same underlying memory. The tensors `x` and `x_`, pointing
+ to the same underlying memory make this transformation safe in the
+ context of EPLB. i.e. It is the same memory and just the view
+ is different.
+ Note: This function handles the "weight_scale" tensors specifically.
+ This could however be generalized to handle similar tensors.
+ """
+ if p.ndim != 3:
+ return p
+ if p.is_contiguous():
+ # Already contiguous. do nothing.
+ return p
+ # p is non-contiguous. We only handle the case where the last 2
+ # dimensions of the scales tensor is transposed. We can handle
+ # other cases when they become relevant.
+ is_transposed_12 = p.stride(1) == 1 and p.stride(2) != 1
+ if "weight_scale" not in name or not is_transposed_12:
+ # do nothing.
+ return p
+
+ # Do not update the layer paramater as the layer's MoE operations would
+ # expect the parameter's tensor to the same shape / stride. Instead,
+ # make a new torch.nn.Parameter that is used just in the context of
+ # EPLB.
+ return torch.nn.Parameter(
+ torch.transpose(p.data, 1, 2), requires_grad=False
+ )
+
weights = list(self.named_parameters())
+ weights = [(name, _maybe_make_contiguous(name, p)) for name, p in weights]
+
assert all(
weight.is_contiguous()
for name, weight in weights
@@ -1381,8 +1588,6 @@ class FusedMoE(CustomOp):
routed_scaling_factor=routed_scaling_factor,
e_score_correction_bias=e_score_correction_bias,
)
- if indices_type is not None:
- topk_ids = topk_ids.to(dtype=indices_type)
elif e_score_correction_bias is not None:
topk_weights, topk_ids = fused_topk_bias(
hidden_states=hidden_states,
@@ -1391,7 +1596,7 @@ class FusedMoE(CustomOp):
topk=top_k,
renormalize=renormalize,
)
- if routed_scaling_factor is not None:
+ if routed_scaling_factor != 1.0:
topk_weights *= routed_scaling_factor
elif custom_routing_function is None:
topk_weights, topk_ids, token_expert_indices = fused_topk(
@@ -1408,8 +1613,6 @@ class FusedMoE(CustomOp):
topk=top_k,
renormalize=renormalize,
)
- if indices_type is not None:
- topk_ids = topk_ids.to(dtype=indices_type)
if enable_eplb:
assert expert_load_view is not None
@@ -1421,9 +1624,11 @@ class FusedMoE(CustomOp):
expert_load_view=expert_load_view,
logical_to_physical_map=logical_to_physical_map,
logical_replica_count=logical_replica_count,
- indices_type=indices_type,
)
+ if (indices_type is not None) and topk_ids.dtype != indices_type:
+ topk_ids = topk_ids.to(dtype=indices_type)
+
assert topk_ids.dtype == indices_type or indices_type is None
# Compute zero expert result if needed
@@ -1694,36 +1899,12 @@ class FusedMoE(CustomOp):
use_chunked_impl = self.use_dp_chunking
- use_shared_experts_stream = (
- has_separate_shared_experts
- and not use_chunked_impl
- and self.shared_experts_stream is not None
- and (
- hidden_states.shape[0]
- <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
+ use_shared_experts_stream, hidden_states_clone = (
+ self._maybe_setup_shared_experts_stream(
+ hidden_states, has_separate_shared_experts, use_chunked_impl
)
)
- if use_shared_experts_stream:
- assert self.shared_experts_stream is not None
-
- # Clone BEFORE switching streams to avoid race condition
- # where routed_expert kernel may mutate hidden_states.
- hidden_states_clone = hidden_states.clone()
-
- # Record that the clone will be used by shared_experts_stream
- # to avoid gc issue from deallocation of hidden_states_clone
- # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
- # NOTE: We dont need shared_output.record_stream(current_stream())
- # because we synch the streams before using shared_output.
- hidden_states_clone.record_stream(self.shared_experts_stream)
-
- # Mark sync start point for the separate shared experts
- # stream here since we want to run in parallel with the
- # router/gate (next op below)
- assert self.shared_experts_stream is not None
- self.shared_experts_stream.wait_stream(current_stream())
-
# If router/gate provided, then apply it here.
# (Note: This code runs only when "overlapped mode" is on to allow
# parallel execution of shared experts with the FusedMoE via
@@ -1752,6 +1933,24 @@ class FusedMoE(CustomOp):
hidden_states_combined, router_logits = get_ep_group().dispatch(
hidden_states, router_logits, self.is_sequence_parallel
)
+ # Run shared experts before matrix multiply.
+ # because matrix multiply maybe modify the hidden_states.
+ if has_separate_shared_experts and not use_shared_experts_stream:
+ assert self.shared_experts is not None
+ shared_output = self.shared_experts(hidden_states)
+
+ # NOTE: Similar with DP, PCP also needs dispatch and combine. For
+ # simplicity, AgRsAll2All was added separately for PCP here. Maybe
+ # we should modify All2AllManager abstract to better support PCP.
+ if self.pcp_size > 1:
+ hidden_states = get_pcp_group().all_gather(
+ hidden_states,
+ dim=0,
+ )
+ router_logits = get_pcp_group().all_gather(
+ router_logits,
+ dim=0,
+ )
# Matrix multiply.
final_hidden_states = self.quant_method.apply(
@@ -1795,8 +1994,6 @@ class FusedMoE(CustomOp):
# conflict with the main stream
shared_output = self.shared_experts(hidden_states_clone)
current_stream().wait_stream(self.shared_experts_stream)
- else:
- shared_output = self.shared_experts(hidden_states)
final_hidden_states = (
shared_output,
@@ -1809,6 +2006,13 @@ class FusedMoE(CustomOp):
def combine_output(states: torch.Tensor) -> torch.Tensor:
if do_naive_dispatch_combine:
states = get_ep_group().combine(states, self.is_sequence_parallel)
+
+ if self.pcp_size > 1:
+ states = get_pcp_group().reduce_scatter(
+ states,
+ dim=0,
+ )
+
return states
if self.shared_experts is not None:
diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py
index 093affe51f503..b2af58cdca887 100644
--- a/vllm/model_executor/layers/fused_moe/modular_kernel.py
+++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py
@@ -10,12 +10,16 @@ from typing import final
import torch
import vllm.envs as envs
+from vllm.config import get_current_vllm_config
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
from vllm.model_executor.layers.fused_moe.utils import (
_resize_cache,
count_expert_num_tokens,
disable_inplace,
)
+from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv
from vllm.v1.worker.ubatching import (
dbo_current_ubatch_id,
@@ -25,6 +29,8 @@ from vllm.v1.worker.ubatching import (
dbo_yield,
)
+logger = init_logger(__name__)
+
#
# This file defines a set of base classes used to make MoE kernels more modular.
# The goal is to be able to utilize different communication mechanisms with
@@ -709,11 +715,13 @@ class FusedMoEModularKernel(torch.nn.Module):
prepare_finalize: FusedMoEPrepareAndFinalize,
fused_experts: FusedMoEPermuteExpertsUnpermute,
shared_experts: torch.nn.Module | None = None,
+ shared_experts_stream: torch.cuda.Stream | None = None,
):
super().__init__()
self.prepare_finalize = prepare_finalize
self.fused_experts = fused_experts
self.shared_experts = shared_experts
+ self.shared_experts_stream = shared_experts_stream
self._post_init_setup()
assert (
@@ -795,6 +803,42 @@ class FusedMoEModularKernel(torch.nn.Module):
buffers = self.shared_buffers[ubatch_idx]
workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)
+ # Force worst-case allocation in profiling run for
+ # "mk.FusedMoEModularKernel.Standard" formats where this is only bounded
+ # by `VLLM_FUSED_MOE_CHUNK_SIZE` and may not be seen during profiling with
+ # DP+EP due to the random token routing.
+ is_profile_run = (
+ is_forward_context_available()
+ and get_forward_context().attn_metadata is None
+ )
+ if is_profile_run and self.fused_experts.supports_chunking():
+ parallel_config = get_current_vllm_config().parallel_config
+ is_dp_ep = (
+ parallel_config.data_parallel_size > 1
+ and parallel_config.enable_expert_parallel
+ )
+ if is_dp_ep:
+ max_workspace_13, max_workspace_2, max_fused_out_shape = (
+ self.fused_experts.workspace_shapes(
+ envs.VLLM_FUSED_MOE_CHUNK_SIZE,
+ N,
+ K,
+ top_k,
+ global_num_experts,
+ local_num_experts,
+ expert_tokens_meta,
+ )
+ )
+ buffers.workspace13.get(
+ max_workspace_13, device=device, dtype=workspace_dtype
+ )
+ buffers.workspace2.get(
+ max_workspace_2, device=device, dtype=workspace_dtype
+ )
+ buffers.fused_out.get(
+ max_fused_out_shape, device=device, dtype=workspace_dtype
+ )
+
# Get intermediate workspace shapes based off the chunked M size.
workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
M_chunk,
@@ -890,6 +934,34 @@ class FusedMoEModularKernel(torch.nn.Module):
expert_num_tokens_cpu=c_expert_num_tokens_cpu,
)
+ def _maybe_setup_shared_experts_stream(
+ self, hidden_states: torch.Tensor
+ ) -> tuple[bool, torch.Tensor | None]:
+ # decide whether to run shared experts on a separate CUDA stream to
+ # overlap with the main fused MoE kernel.
+ use_shared_experts_stream = (
+ self.shared_experts is not None
+ and self.shared_experts_stream is not None
+ and hidden_states.is_cuda
+ and (
+ hidden_states.shape[0]
+ <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
+ )
+ )
+
+ hidden_states_clone: torch.Tensor | None = None
+ if use_shared_experts_stream and self.shared_experts_stream is not None:
+ # TODO: Optimize this (complicated)
+ # Note: this clone adds overhead but is required
+ # for correctness with multiple CUDA streams and CUDA graph capture.
+ hidden_states_clone = hidden_states.clone()
+ # record that the clone will be used by the separate stream so its
+ # lifetime is correctly tracked.
+ hidden_states_clone.record_stream(self.shared_experts_stream)
+ self.shared_experts_stream.wait_stream(torch.cuda.current_stream())
+
+ return use_shared_experts_stream, hidden_states_clone
+
def _prepare(
self,
hidden_states: torch.Tensor,
@@ -1077,12 +1149,30 @@ class FusedMoEModularKernel(torch.nn.Module):
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
apply_router_weight_on_input: bool,
+ hidden_states_clone: torch.Tensor | None = None,
+ use_shared_experts_stream: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""
The _finalize method is a wrapper around self.prepare_finalize.finalize
that handles DBO, async and shared expert overlap.
"""
- shared_output: torch.Tensor | None = None
+
+ def maybe_run_shared_experts() -> torch.Tensor | None:
+ if self.shared_experts is None:
+ return None
+
+ if (
+ not use_shared_experts_stream
+ or self.shared_experts_stream is not None
+ and (not hidden_states.is_cuda or not torch.cuda.is_available())
+ ):
+ # fall back to running on the current stream
+ return self.shared_experts(hidden_states)
+
+ assert hidden_states_clone is not None
+ # launch shared experts on the dedicated stream.
+ with torch.cuda.stream(self.shared_experts_stream):
+ return self.shared_experts(hidden_states_clone)
if not self.prepare_finalize.supports_async():
assert not dbo_enabled()
@@ -1095,8 +1185,7 @@ class FusedMoEModularKernel(torch.nn.Module):
apply_router_weight_on_input,
self.fused_experts.finalize_weight_and_reduce_impl(),
)
- if self.shared_experts is not None:
- shared_output = self.shared_experts(hidden_states)
+ shared_output = maybe_run_shared_experts()
else:
finalize_ret = self.prepare_finalize.finalize_async(
output,
@@ -1107,8 +1196,7 @@ class FusedMoEModularKernel(torch.nn.Module):
self.fused_experts.finalize_weight_and_reduce_impl(),
)
- if self.shared_experts is not None:
- shared_output = self.shared_experts(hidden_states)
+ shared_output = maybe_run_shared_experts()
# TODO(lucas): refactor this in the alternative schedules followup
# currently unpack if we have hook + receiver pair or just
@@ -1131,12 +1219,28 @@ class FusedMoEModularKernel(torch.nn.Module):
receiver()
+ self._wait_for_shared_experts_stream(hidden_states, use_shared_experts_stream)
+
if self.shared_experts is None:
return output
else:
assert shared_output is not None
return shared_output, output
+ def _wait_for_shared_experts_stream(
+ self, hidden_states: torch.Tensor, use_shared_experts_stream: bool
+ ) -> None:
+ # ensure that any work enqueued on the shared_experts_stream is
+ # completed before the shared_output tensor is consumed
+ if (
+ self.shared_experts is not None
+ and use_shared_experts_stream
+ and self.shared_experts_stream is not None
+ and hidden_states.is_cuda
+ and current_platform.is_cuda()
+ ):
+ torch.cuda.current_stream().wait_stream(self.shared_experts_stream)
+
def forward(
self,
hidden_states: torch.Tensor,
@@ -1183,6 +1287,10 @@ class FusedMoEModularKernel(torch.nn.Module):
else:
output = torch.zeros_like(hidden_states)
+ use_shared_experts_stream, hidden_states_clone = (
+ self._maybe_setup_shared_experts_stream(hidden_states)
+ )
+
local_num_experts = w1.size(0)
if global_num_experts == -1:
global_num_experts = local_num_experts
@@ -1219,4 +1327,6 @@ class FusedMoEModularKernel(torch.nn.Module):
topk_weights,
topk_ids,
apply_router_weight_on_input,
+ hidden_states_clone=hidden_states_clone,
+ use_shared_experts_stream=use_shared_experts_stream,
)
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index 9bb976fb9ec93..e27e2eb32da0f 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -45,7 +45,8 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
assert topk == 1, (
"apply_router_weight_on_input is only implemented for topk=1"
)
- a1.mul_(topk_weights.to(a1.dtype))
+ # Note: do not use inplace for shared experts overlap
+ a1 = a1 * topk_weights.to(a1.dtype)
a1q, a1q_scale = moe_kernel_quantize_input(
a1,
diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
index ce56887f1c26d..63b0e6f573d65 100644
--- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
@@ -108,11 +108,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
def allow_inplace(self) -> bool:
return True
- def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:
+ def maybe_make_prepare_finalize(
+ self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+ ) -> FusedMoEPrepareAndFinalize | None:
if self.rocm_aiter_moe_enabled:
return None
else:
- return super().maybe_make_prepare_finalize()
+ return super().maybe_make_prepare_finalize(routing_tables)
def select_gemm_impl(
self,
@@ -260,7 +263,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
layer.w2_weight.copy_(packed_w2_weight)
layer.cpu_fused_moe = cpu_fused_moe.SGLFusedMOE(layer)
else:
- layer.cpu_fused_moe = cpu_fused_moe.IPEXFusedMOE(layer)
+ layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
else:
layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
diff --git a/vllm/model_executor/layers/kda.py b/vllm/model_executor/layers/kda.py
index 2e7500bac7188..27cc3884517f9 100644
--- a/vllm/model_executor/layers/kda.py
+++ b/vllm/model_executor/layers/kda.py
@@ -5,7 +5,6 @@ import torch
from einops import rearrange
from torch import nn
-from vllm.attention import AttentionBackend
from vllm.attention.backends.abstract import AttentionMetadata
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
from vllm.distributed import (
@@ -83,12 +82,7 @@ direct_register_custom_op(
class KimiDeltaAttention(nn.Module, MambaBase):
@property
def mamba_type(self) -> str:
- return "linear_attention"
-
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.gdn_attn import GDNAttentionBackend
-
- return GDNAttentionBackend
+ return "gdn_attention"
def get_state_dtype(
self,
diff --git a/vllm/model_executor/layers/lightning_attn.py b/vllm/model_executor/layers/lightning_attn.py
index 99853680eac6c..ffccdc12241cb 100644
--- a/vllm/model_executor/layers/lightning_attn.py
+++ b/vllm/model_executor/layers/lightning_attn.py
@@ -198,7 +198,7 @@ def _fwd_kv_parallel(
)
# Load the decay factors for the current head and block
- k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)[None, :]
+ k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)
kv_index = tl.arange(0, CBLOCK)
@@ -228,6 +228,12 @@ def _fwd_kv_parallel(
# Load decay factor and compute weighted key-value outer product
k_decay = tl.load(k_decay_ptr)
+
+ # NOTE: Need to add the extra dim here due to AMD MLIR lowering error.
+ # Please don't move it back until issue is resolved.
+ # Issue: https://github.com/ROCm/triton/issues/907
+ k_decay = k_decay[None, :]
+
kv += tl.dot(k_trans * k_decay, v)
# Move to the next sub-block
diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py
index e68b09b4d81f5..aa919d6fdc35c 100644
--- a/vllm/model_executor/layers/mamba/abstract.py
+++ b/vllm/model_executor/layers/mamba/abstract.py
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
import torch
+from vllm.attention.selector import get_mamba_attn_backend
from vllm.config import VllmConfig
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
@@ -38,11 +39,6 @@ class MambaBase(AttentionLayerBase):
def mamba_type(self) -> str:
pass
- @abstractmethod
- def get_attn_backend(self) -> type["AttentionBackend"]:
- """Get the attention backend class for this Mamba layer."""
- pass
-
@abstractmethod
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
pass
@@ -69,3 +65,7 @@ class MambaBase(AttentionLayerBase):
else 0
),
)
+
+ def get_attn_backend(self) -> type["AttentionBackend"]:
+ """Get the attention backend class for this Mamba layer."""
+ return get_mamba_attn_backend(self.mamba_type)
diff --git a/vllm/model_executor/layers/mamba/linear_attn.py b/vllm/model_executor/layers/mamba/linear_attn.py
index 0a2742ff49a44..d85b3e61c5d61 100644
--- a/vllm/model_executor/layers/mamba/linear_attn.py
+++ b/vllm/model_executor/layers/mamba/linear_attn.py
@@ -2,12 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionBackend
-
-from typing import TYPE_CHECKING
import torch
import torch.nn.functional as F
@@ -37,9 +31,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.utils.torch_utils import direct_register_custom_op
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
-if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionBackend
-
class MiniMaxText01RMSNormTP(CustomOp):
name = "MiniMaxText01RMSNormTP"
@@ -123,11 +114,6 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
def mamba_type(self) -> str:
return "linear_attention"
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
-
- return LinearAttentionBackend
-
def get_state_dtype(self) -> tuple[torch.dtype]:
assert self.model_config is not None
assert self.cache_config is not None
diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py
index b6345b8af7f0a..90e520e244416 100644
--- a/vllm/model_executor/layers/mamba/mamba_mixer.py
+++ b/vllm/model_executor/layers/mamba/mamba_mixer.py
@@ -1,10 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING, NamedTuple
-
-if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionBackend
+from typing import NamedTuple
import torch
from torch import nn
@@ -452,11 +449,6 @@ class MambaMixer(MambaBase, CustomOp):
def mamba_type(self) -> str:
return "mamba1"
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
-
- return Mamba1AttentionBackend
-
def _time_proj_bias(self) -> torch.Tensor | None:
if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None:
return self.dt_proj.bias.float()
diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py
index fb45afa33dad6..0ea5805305eda 100644
--- a/vllm/model_executor/layers/mamba/mamba_mixer2.py
+++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py
@@ -1,10 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionBackend
import torch
from torch import nn
@@ -426,6 +422,10 @@ class MambaMixer2(MambaBase, CustomOp):
# `ColumnParallelLinear` and `MergedColumnParallelLinear`,
# and `set_weight_attrs` doesn't allow to override it
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
+ conv_weights = self.conv1d.weight.view(
+ self.conv1d.weight.size(0), self.conv1d.weight.size(2)
+ )
+ self.register_buffer("conv_weights", conv_weights, persistent=False)
# - these are TPed by heads to reduce the size of the
# temporal shape
@@ -459,6 +459,17 @@ class MambaMixer2(MambaBase, CustomOp):
intermediate_size, n_groups, self.use_rms_norm, eps=rms_norm_eps
)
+ # - get hidden_states, B and C after depthwise convolution.
+ self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split(
+ hidden_states_B_C,
+ [
+ self.intermediate_size // self.tp_size,
+ self.groups_ssm_state_size // self.tp_size,
+ self.groups_ssm_state_size // self.tp_size,
+ ],
+ dim=-1,
+ )
+
compilation_config = get_current_vllm_config().compilation_config
if prefix in compilation_config.static_forward_context:
raise ValueError(f"Duplicate layer name: {prefix}")
@@ -470,10 +481,24 @@ class MambaMixer2(MambaBase, CustomOp):
self.cache_config = cache_config
self.prefix = prefix
+ # Pre-compute sizes for forward pass
+ self.tped_intermediate_size = self.intermediate_size // self.tp_size
+ self.tped_conv_size = self.conv_dim // self.tp_size
+ self.tped_dt_size = self.num_heads // self.tp_size
+
+ self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split(
+ hidden_states_B_C,
+ [
+ self.tped_intermediate_size,
+ self.groups_ssm_state_size // self.tp_size,
+ self.groups_ssm_state_size // self.tp_size,
+ ],
+ dim=-1,
+ )
+
def forward_native(
self,
hidden_states: torch.Tensor,
- output: torch.Tensor,
mup_vector: torch.Tensor | None = None,
):
pass
@@ -481,22 +506,55 @@ class MambaMixer2(MambaBase, CustomOp):
def forward(
self,
hidden_states: torch.Tensor,
- output: torch.Tensor,
mup_vector: torch.Tensor | None = None,
):
- torch.ops.vllm.mamba_mixer2(
- hidden_states,
- output,
- self.prefix,
- mup_vector,
+ # 1. Gated MLP's linear projection
+ projected_states, _ = self.in_proj(hidden_states)
+ if mup_vector is not None:
+ projected_states = projected_states * mup_vector
+
+ # 2. Prepare inputs for conv + SSM
+ ssm_output = torch.empty(
+ [
+ hidden_states.shape[0],
+ (self.num_heads // self.tp_size) * self.head_dim,
+ ],
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
)
- def forward_cuda(
+ # 3. conv + SSM
+ # (split `projected_states` into hidden_states_B_C, dt in the custom op to
+ # ensure it is not treated as an intermediate tensor by torch compile)
+ torch.ops.vllm.mamba_mixer2(
+ projected_states,
+ ssm_output,
+ self.prefix,
+ )
+
+ # 4. gated MLP
+ # GatedRMSNorm internally applying SiLU to the gate
+ # SiLU is applied internally before normalization, unlike standard
+ # norm usage
+ gate = projected_states[..., : self.tped_intermediate_size]
+ hidden_states = self.norm(ssm_output, gate)
+
+ # 5. Final linear projection
+ output, _ = self.out_proj(hidden_states)
+
+ return output
+
+ def conv_ssm_forward(
self,
- hidden_states: torch.Tensor,
+ projected_states: torch.Tensor,
output: torch.Tensor,
- mup_vector: torch.Tensor | None = None,
):
+ hidden_states_B_C, dt = torch.split(
+ projected_states[..., self.tped_intermediate_size :],
+ [self.tped_conv_size, self.tped_dt_size],
+ dim=-1,
+ )
+
forward_context = get_forward_context()
# attn_metadata contains metadata necessary for the mamba2 triton
# kernels to operate in continuous batching and in chunked prefill
@@ -524,48 +582,14 @@ class MambaMixer2(MambaBase, CustomOp):
cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
last_chunk_indices_p = attn_metadata.last_chunk_indices_p
- # 1. Gated MLP's linear projection
- projected_states, _ = self.in_proj(hidden_states)
-
- if mup_vector is not None:
- projected_states = projected_states * mup_vector
-
- gate, hidden_states_B_C, dt = torch.split(
- projected_states,
- [
- self.intermediate_size // self.tp_size,
- self.conv_dim // self.tp_size,
- self.num_heads // self.tp_size,
- ],
- dim=-1,
- )
-
- conv_weights = self.conv1d.weight.view(
- self.conv1d.weight.size(0), self.conv1d.weight.size(2)
- )
-
- # - get hidden_states, B and C after depthwise convolution.
- split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split(
- hidden_states_B_C,
- [
- self.intermediate_size // self.tp_size,
- self.groups_ssm_state_size // self.tp_size,
- self.groups_ssm_state_size // self.tp_size,
- ],
- dim=-1,
- )
-
if attn_metadata is None:
# profile run
hidden_states_B_C = (
hidden_states_B_C.transpose(0, 1).clone().transpose(0, 1)
).contiguous()
- hidden_states, _B, _C = split_hidden_states_B_C_fn(hidden_states_B_C)
- hidden_states = self.norm(hidden_states, gate)
- out, _ = self.out_proj(hidden_states)
- return out
+ hidden_states, _B, _C = self.split_hidden_states_B_C_fn(hidden_states_B_C)
+ return hidden_states
- # NOTE: V0 put prefill before decode, v1 puts decode before prefill
num_prefills = attn_metadata.num_prefills # request count
num_decodes = attn_metadata.num_decode_tokens # token count (=request)
num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
@@ -622,18 +646,8 @@ class MambaMixer2(MambaBase, CustomOp):
block_idx_first_scheduled_token_p = None
num_computed_tokens_p = None
- # Preallocate output tensor to avoid memcpy cost for merging prefill
- # and decode outputs
- preallocated_ssm_out = torch.empty(
- [
- num_prefill_tokens + num_decodes,
- (self.num_heads // self.tp_size) * self.head_dim,
- ],
- dtype=hidden_states.dtype,
- device=hidden_states.device,
- )
preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
- preallocated_ssm_out,
+ output[:num_actual_tokens],
[num_decodes, num_prefill_tokens],
dim=0,
)
@@ -658,7 +672,7 @@ class MambaMixer2(MambaBase, CustomOp):
) # this is the form that causal-conv see
hidden_states_B_C_p = causal_conv1d_fn(
x,
- conv_weights,
+ self.conv_weights,
self.conv1d.bias,
activation=self.activation,
conv_states=conv_state,
@@ -673,7 +687,9 @@ class MambaMixer2(MambaBase, CustomOp):
query_start_loc=query_start_loc_p,
).transpose(0, 1)[:num_prefill_tokens]
- hidden_states_p, B_p, C_p = split_hidden_states_B_C_fn(hidden_states_B_C_p)
+ hidden_states_p, B_p, C_p = self.split_hidden_states_B_C_fn(
+ hidden_states_B_C_p
+ )
# 3. State Space Model sequence transformation
initial_states = None
@@ -815,7 +831,7 @@ class MambaMixer2(MambaBase, CustomOp):
hidden_states_B_C_d = causal_conv1d_update(
hidden_states_B_C_d,
conv_state,
- conv_weights,
+ self.conv_weights,
self.conv1d.bias,
self.activation,
conv_state_indices=state_indices_tensor_d,
@@ -823,7 +839,9 @@ class MambaMixer2(MambaBase, CustomOp):
initial_state_idx=block_idx_last_computed_token_d,
)
- hidden_states_d, B_d, C_d = split_hidden_states_B_C_fn(hidden_states_B_C_d)
+ hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn(
+ hidden_states_B_C_d
+ )
# 3. State Space Model sequence transformation
n_groups = self.n_groups // self.tp_size
@@ -861,15 +879,6 @@ class MambaMixer2(MambaBase, CustomOp):
out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
)
- # 4. gated MLP
- # GatedRMSNorm internally applying SiLU to the gate
- # SiLU is applied internally before normalization, unlike standard
- # norm usage
- hidden_states = self.norm(preallocated_ssm_out, gate[:num_actual_tokens])
-
- # 5. Final linear projection
- output[:num_actual_tokens], _ = self.out_proj(hidden_states)
-
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
assert self.model_config is not None
assert self.cache_config is not None
@@ -894,28 +903,21 @@ class MambaMixer2(MambaBase, CustomOp):
def mamba_type(self) -> str:
return "mamba2"
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
-
- return Mamba2AttentionBackend
-
def mamba_mixer2(
- hidden_states: torch.Tensor,
+ projected_states: torch.Tensor,
output: torch.Tensor,
layer_name: str,
- mup_vector: torch.Tensor | None = None,
) -> None:
forward_context: ForwardContext = get_forward_context()
self = forward_context.no_compile_layers[layer_name]
- self.forward_cuda(hidden_states=hidden_states, output=output, mup_vector=mup_vector)
+ self.conv_ssm_forward(projected_states=projected_states, output=output)
def mamba_mixer2_fake(
- hidden_states: torch.Tensor,
+ projected_states: torch.Tensor,
output: torch.Tensor,
layer_name: str,
- mup_vector: torch.Tensor | None = None,
) -> None:
return
diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py
index 04efa8a8b3734..0bbad17d7ebc7 100644
--- a/vllm/model_executor/layers/mamba/short_conv.py
+++ b/vllm/model_executor/layers/mamba/short_conv.py
@@ -1,10 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionBackend
import torch
@@ -232,11 +228,6 @@ class ShortConv(MambaBase, CustomOp):
def mamba_type(self) -> str:
return "short_conv"
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
-
- return ShortConvAttentionBackend
-
def short_conv(
hidden_states: torch.Tensor,
diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py
index c4c44b83ae6bf..6ebfa47a9dc3f 100644
--- a/vllm/model_executor/layers/mla.py
+++ b/vllm/model_executor/layers/mla.py
@@ -24,6 +24,7 @@ class MLAModules:
q_b_proj: torch.nn.Module | None
q_proj: torch.nn.Module | None
indexer: torch.nn.Module | None
+ indexer_rotary_emb: torch.nn.Module | None
is_sparse: bool
topk_indices_buffer: torch.Tensor | None
@@ -80,6 +81,7 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
self.rotary_emb = mla_modules.rotary_emb
self.o_proj = mla_modules.o_proj
self.indexer = mla_modules.indexer
+ self.indexer_rope_emb = mla_modules.indexer_rotary_emb
self.is_sparse = mla_modules.is_sparse
if self.indexer is not None:
@@ -153,7 +155,9 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
)
if self.indexer and self.is_sparse:
- _topk_indices = self.indexer(hidden_states, q_c, positions, self.rotary_emb)
+ _topk_indices = self.indexer(
+ hidden_states, q_c, positions, self.indexer_rope_emb
+ )
attn_out = self.mla_attn(
q,
diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py
index bb42b10f87186..18aaae394f935 100644
--- a/vllm/model_executor/layers/quantization/__init__.py
+++ b/vllm/model_executor/layers/quantization/__init__.py
@@ -38,6 +38,8 @@ QuantizationMethods = Literal[
"inc",
"mxfp4",
"petit_nvfp4",
+ "cpu_gptq",
+ "cpu_awq",
]
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
@@ -107,6 +109,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
from .compressed_tensors.compressed_tensors import (
CompressedTensorsConfig,
)
+ from .cpu_wna16 import CPUAWQConfig, CPUGPTQConfig
from .deepspeedfp import DeepSpeedFPConfig
from .experts_int8 import ExpertsInt8Config
from .fbgemm_fp8 import FBGEMMFp8Config
@@ -159,6 +162,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
"inc": INCConfig,
"mxfp4": Mxfp4Config,
"petit_nvfp4": PetitNvFp4Config,
+ "cpu_gptq": CPUGPTQConfig,
+ "cpu_awq": CPUAWQConfig,
}
# Update the `method_to_config` with customized quantization methods.
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
index 06ee96d55419c..ad547dd409822 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
@@ -8,6 +8,7 @@ from enum import Enum
import torch
from compressed_tensors import CompressionFormat
from compressed_tensors.quantization import ActivationOrdering, QuantizationStrategy
+from torch.nn.parameter import Parameter
import vllm.envs as envs
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -50,9 +51,15 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
build_flashinfer_fp4_cutlass_moe_prepare_finalize,
+ flashinfer_trtllm_fp4_moe,
+ prepare_static_weights_for_trtllm_fp4_moe,
reorder_w1w3_to_w3w1,
select_nvfp4_gemm_impl,
)
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+ FlashinferMoeBackend,
+ get_flashinfer_moe_backend,
+)
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
expert_weight_is_col_major,
requant_weight_ue8m0_inplace,
@@ -193,6 +200,13 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
self.allow_flashinfer = _nvfp4.allow_flashinfer
self.use_marlin = _nvfp4.use_marlin
self.group_size = 16
+ self.flashinfer_moe_backend = None
+ if self.allow_flashinfer:
+ self.flashinfer_moe_backend = get_flashinfer_moe_backend()
+ logger.info_once(
+ f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
+ " for CompressedTensorsW4A4MoeMethod."
+ )
def create_weights(
self,
@@ -344,21 +358,20 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
if self.use_marlin:
prepare_moe_fp4_layer_for_marlin(layer)
return
-
- # swizzle weight scales
- layer.w13_weight_scale = torch.nn.Parameter(
- swizzle_blockscale(layer.w13_weight_scale), requires_grad=False
- )
-
- layer.w2_weight_scale = torch.nn.Parameter(
- swizzle_blockscale(layer.w2_weight_scale), requires_grad=False
- )
-
# w13
- w13_input_global_scale = layer.w13_input_global_scale.max(dim=1).values.to(
- torch.float32
- )
-
+ if (
+ self.allow_flashinfer
+ and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
+ w13_input_global_scale = (
+ layer.w13_input_global_scale.min()
+ .to(torch.float32)
+ .expand(layer.num_experts)
+ )
+ else:
+ w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to(
+ torch.float32
+ )
layer.g1_alphas = torch.nn.Parameter(
((1 / w13_input_global_scale) * layer.w13_weight_scale_2),
requires_grad=False,
@@ -369,22 +382,95 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
)
# w2
+ if (
+ self.allow_flashinfer
+ and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
+ w2_input_global_scale = (
+ layer.w2_input_global_scale.min()
+ .to(torch.float32)
+ .expand(layer.num_experts)
+ )
+ else:
+ w2_input_global_scale = layer.w2_input_global_scale
+
layer.g2_alphas = torch.nn.Parameter(
- ((1 / layer.w2_input_global_scale) * layer.w2_weight_scale_2).to(
- torch.float32
- ),
+ ((1 / w2_input_global_scale) * layer.w2_weight_scale_2).to(torch.float32),
requires_grad=False,
)
layer.w2_input_scale_quant = torch.nn.Parameter(
- (layer.w2_input_global_scale), requires_grad=False
+ (w2_input_global_scale), requires_grad=False
)
- def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:
- if self.use_marlin:
+ # TensorRT-LLM specific processing
+ if (
+ self.allow_flashinfer
+ and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
+ # Prepare static weights for TRT-LLM kernel
+ # alternate: prepare_static_weight_layouts_for_trtllm_moe
+ (
+ gemm1_weights_fp4_shuffled,
+ gemm1_scales_fp4_shuffled,
+ gemm2_weights_fp4_shuffled,
+ gemm2_scales_fp4_shuffled,
+ ) = prepare_static_weights_for_trtllm_fp4_moe(
+ layer.w13_weight,
+ layer.w2_weight,
+ layer.w13_weight_scale,
+ layer.w2_weight_scale,
+ layer.w2_weight.size(-2), # hidden_size
+ layer.w13_weight.size(-2) // 2, # intermediate_size
+ layer.w13_weight.size(0), # num_experts
+ )
+ logger.debug_once("Finished shuffling weights for TRT-LLM MOE")
+
+ layer.gemm1_weights_fp4_shuffled = Parameter(
+ gemm1_weights_fp4_shuffled, requires_grad=False
+ )
+ layer.gemm2_weights_fp4_shuffled = Parameter(
+ gemm2_weights_fp4_shuffled, requires_grad=False
+ )
+ layer.gemm1_scales_fp4_shuffled = Parameter(
+ gemm1_scales_fp4_shuffled, requires_grad=False
+ )
+ layer.gemm2_scales_fp4_shuffled = Parameter(
+ gemm2_scales_fp4_shuffled, requires_grad=False
+ )
+
+ # Additional parameter needed for TRT-LLM
+ layer.g1_scale_c = Parameter(
+ (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),
+ requires_grad=False,
+ )
+
+ # Clean up weights that won't be used by TRT-LLM
+ del layer.w2_weight
+ del layer.w2_weight_scale
+ del layer.w13_weight
+ del layer.w13_weight_scale
+ else:
+ # swizzle weight scales
+ layer.w13_weight_scale = torch.nn.Parameter(
+ swizzle_blockscale(layer.w13_weight_scale), requires_grad=False
+ )
+
+ layer.w2_weight_scale = torch.nn.Parameter(
+ swizzle_blockscale(layer.w2_weight_scale), requires_grad=False
+ )
+
+ def maybe_make_prepare_finalize(
+ self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+ ) -> mk.FusedMoEPrepareAndFinalize | None:
+ if self.use_marlin or (
+ self.allow_flashinfer
+ and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
return None
elif not self.allow_flashinfer:
- return super().maybe_make_prepare_finalize()
+ return super().maybe_make_prepare_finalize(routing_tables)
prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(self.moe)
logger.debug_once("%s", prepare_finalize.__class__.__name__)
@@ -408,7 +494,10 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> FusedMoEQuantConfig | None:
- if self.use_marlin:
+ if (
+ self.use_marlin
+ or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
return None
return nvfp4_moe_quant_config(
@@ -449,6 +538,22 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
)
assert activation == "silu", "Only SiLU activation is supported."
+ if (
+ self.allow_flashinfer
+ and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
+ return flashinfer_trtllm_fp4_moe(
+ layer=layer,
+ x=x,
+ router_logits=router_logits,
+ top_k=top_k,
+ global_num_experts=global_num_experts,
+ num_expert_group=num_expert_group,
+ topk_group=topk_group,
+ custom_routing_function=custom_routing_function,
+ e_score_correction_bias=e_score_correction_bias,
+ )
+
topk_weights, topk_ids, _ = FusedMoE.select_experts(
hidden_states=x,
router_logits=router_logits,
@@ -890,11 +995,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
layer.w2_weight_scale
)
- def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:
+ def maybe_make_prepare_finalize(
+ self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+ ) -> mk.FusedMoEPrepareAndFinalize | None:
if self.use_marlin or self.rocm_aiter_moe_enabled:
return None
else:
- return super().maybe_make_prepare_finalize()
+ return super().maybe_make_prepare_finalize(routing_tables)
def select_gemm_impl(
self,
@@ -1915,9 +2023,20 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
logical_replica_count: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
if enable_eplb:
- raise NotImplementedError(
- "EPLB not supported for `CompressedTensorsWNA16MoEMethod` yet."
- )
+ if expert_load_view is None:
+ raise ValueError("enable_eplb=True requiere expert_load_view != None")
+ if logical_to_physical_map is None:
+ raise ValueError(
+ "enable_eplb=True requiere logical_to_physical_map != None"
+ )
+ if logical_replica_count is None:
+ raise ValueError(
+ "enable_eplb=True requiere logical_replica_count != None"
+ )
+ if not isinstance(layer, FusedMoE):
+ raise TypeError(
+ "EPLB is only supported when `layer` is a instance of FusedMoE."
+ )
from vllm.model_executor.layers.fused_moe import fused_experts
@@ -1934,6 +2053,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
routed_scaling_factor=routed_scaling_factor,
e_score_correction_bias=e_score_correction_bias,
indices_type=self.topk_indices_dtype,
+ num_fused_shared_experts=getattr(layer, "num_fused_shared_experts", 0),
+ enable_eplb=enable_eplb,
+ expert_map=expert_map,
+ expert_load_view=expert_load_view,
+ logical_to_physical_map=logical_to_physical_map,
+ logical_replica_count=logical_replica_count,
)
return fused_experts(
@@ -1950,6 +2075,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
quant_config=self.moe_quant_config,
)
+ @property
+ def supports_eplb(self) -> bool:
+ return True
+
class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
"""
diff --git a/vllm/model_executor/layers/quantization/cpu_wna16.py b/vllm/model_executor/layers/quantization/cpu_wna16.py
new file mode 100644
index 0000000000000..bf643f55f1b9a
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/cpu_wna16.py
@@ -0,0 +1,625 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any, Optional
+
+import torch
+from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
+
+from vllm._custom_ops import (
+ cpu_gemm_wna16,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import (
+ LinearBase,
+ LinearMethodBase,
+ UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization import QuantizationMethods
+from vllm.model_executor.layers.quantization.base_config import (
+ QuantizationConfig,
+ QuantizeMethodBase,
+)
+from vllm.model_executor.layers.quantization.utils.gptq_utils import (
+ get_linear_quant_method,
+)
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+ marlin_repeat_scales_on_all_ranks,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+ is_layer_skipped,
+ pack_cols,
+ unpack_cols,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.model_executor.parameter import (
+ ChannelQuantScaleParameter,
+ GroupQuantScaleParameter,
+ PackedColumnParameter,
+ PackedvLLMParameter,
+ RowvLLMParameter,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
+from vllm.transformers_utils.config import get_safetensors_params_metadata
+from vllm.utils.collection_utils import is_list_of
+
+logger = init_logger(__name__)
+
+
+class CPUGPTQConfig(QuantizationConfig):
+ """Config class for CPU GPTQ quant"""
+
+ def __init__(
+ self,
+ weight_bits: int,
+ group_size: int,
+ desc_act: bool,
+ is_sym: bool,
+ lm_head_quantized: bool,
+ dynamic: dict[str, dict[str, int | bool]],
+ full_config: dict[str, Any],
+ modules_in_block_to_quantize: list[str] | None = None,
+ ) -> None:
+ super().__init__()
+ if desc_act and group_size == -1:
+ # In this case, act_order == True is the same as act_order == False
+ # (since we have only one group per output channel)
+ desc_act = False
+
+ # GPTQModel use `dynamic` config property to allow per module
+ # quantization config so each module can be individually optimized.
+ # Format is dict[str, dict] where key is a regex string that can
+ # perform both positive ("+:" prefixed) or negative ("-:" prefixed)
+ # matching of a module.
+ # Default to positive match, override base quant config mode, if no
+ # prefix is used. Value is in dict format of field key and override
+ # value.
+ # Negative matching will skip quantization init for this module
+ # entirely:
+ # non-quantized inference. More details and quantization examples can be
+ # found at: https://github.com/ModelCloud/GPTQModel
+ # Example:
+ # # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9
+ # # last 1/4 of the layers 16-21 has 8bit and group_size 64
+ # dynamic = {
+ # #`.*\.` matches the layers_node prefix
+ # # positive match layer 10-15
+ # r"+:.*\.(?:1[0-5])\..*": {"bits": 8,},
+ # # positive match layer 16-21
+ # r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,},
+ # r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers
+ # }
+ assert weight_bits == 4
+ self.dynamic = dynamic
+ self.weight_bits = weight_bits
+ self.is_sym = is_sym
+ self.pack_factor = 32 // weight_bits # packed into int32
+ self.group_size = group_size
+ self.desc_act = desc_act
+ self.lm_head_quantized = lm_head_quantized
+ self.full_config = full_config
+ self.modules_in_block_to_quantize = modules_in_block_to_quantize or []
+
+ def __repr__(self) -> str:
+ return (
+ f"CPUWNA16Config("
+ f"group_size={self.group_size}, "
+ f"desc_act={self.desc_act}, "
+ f"lm_head_quantized={self.lm_head_quantized}, "
+ f"dynamic={self.dynamic}, "
+ f"modules_in_block_to_quantize={self.modules_in_block_to_quantize})"
+ )
+
+ @classmethod
+ def get_name(cls) -> QuantizationMethods:
+ return "cpu_gptq"
+
+ @classmethod
+ def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+ return [torch.half, torch.bfloat16]
+
+ @classmethod
+ def get_min_capability(cls) -> int:
+ return -1
+
+ @classmethod
+ def get_config_filenames(cls) -> list[str]:
+ return ["quantize_config.json"]
+
+ @classmethod
+ def from_config(cls, config: dict[str, Any]) -> "CPUGPTQConfig":
+ weight_bits = cls.get_from_keys(config, ["bits"])
+ desc_act = cls.get_from_keys_or(config, ["desc_act"], default=False)
+ dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})
+ group_size = cls.get_from_keys(config, ["group_size"])
+ is_sym = cls.get_from_keys(config, ["sym"])
+ lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
+ modules_in_block_to_quantize = cls.get_from_keys_or(
+ config, ["modules_in_block_to_quantize"], default=None
+ )
+ return cls(
+ weight_bits,
+ group_size,
+ desc_act,
+ is_sym,
+ lm_head_quantized,
+ dynamic,
+ config,
+ modules_in_block_to_quantize,
+ )
+
+ @classmethod
+ def override_quantization_method(
+ cls, hf_quant_cfg, user_quant
+ ) -> QuantizationMethods | None:
+ quant_method = hf_quant_cfg.get("quant_method", "").lower()
+ if current_platform.is_cpu() and (quant_method == "gptq"):
+ return cls.get_name()
+ return None
+
+ def get_quant_method(
+ self, layer: torch.nn.Module, prefix: str
+ ) -> Optional["QuantizeMethodBase"]:
+ return get_linear_quant_method(self, layer, prefix, CPUGPTQLinearMethod) # type: ignore
+
+ def apply_vllm_mapper(self, hf_to_vllm_mapper):
+ if self.modules_in_block_to_quantize is not None:
+ self.modules_in_block_to_quantize = hf_to_vllm_mapper.apply_list(
+ self.modules_in_block_to_quantize
+ )
+
+ def maybe_update_config(self, model_name: str, revision: str | None = None):
+ if self.modules_in_block_to_quantize:
+ if is_list_of(self.modules_in_block_to_quantize, list):
+ # original modules_in_block_to_quantize: list[list[str]]
+ # flatten original modules_in_block_to_quantize
+ self.modules_in_block_to_quantize = [
+ item
+ for sublist in self.modules_in_block_to_quantize
+ for item in sublist
+ ]
+ return
+
+ unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+ metadata = get_safetensors_params_metadata(model_name, revision=revision)
+ quant_layers: set[str] = {
+ param_name.rsplit(".", 1)[0]
+ for param_name, info in metadata.items()
+ if (dtype := info.get("dtype", None))
+ and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
+ }
+ self.modules_in_block_to_quantize = list(quant_layers)
+
+
+class CPUGPTQLinearMethod(LinearMethodBase):
+ """Linear method for GPTQ on CPU.
+
+ Args:
+ quant_config: The CPUWNA16 quantization config.
+ """
+
+ def __init__(self, quant_config: CPUGPTQConfig) -> None:
+ self.quant_config = quant_config
+ assert self.quant_config.is_sym, "GPTQ asym quant is not supported on CPU"
+
+ def create_weights(
+ self,
+ layer: torch.nn.Module,
+ input_size_per_partition: int,
+ output_partition_sizes: list[int],
+ input_size: int,
+ output_size: int,
+ params_dtype: torch.dtype,
+ **extra_weight_attrs,
+ ) -> None:
+ output_size_per_partition = sum(output_partition_sizes)
+ assert output_size_per_partition * self.quant_config.weight_bits % 32 == 0
+ assert output_size_per_partition % 32 == 0
+ assert input_size_per_partition % 32 == 0
+
+ is_row_parallel = input_size != input_size_per_partition
+ weight_loader = extra_weight_attrs.get("weight_loader")
+
+ # Normalize group_size
+ if self.quant_config.group_size != -1:
+ group_size = self.quant_config.group_size
+ else:
+ group_size = input_size
+
+ # Determine sharding
+ if marlin_repeat_scales_on_all_ranks(
+ self.quant_config.desc_act, self.quant_config.group_size, is_row_parallel
+ ):
+ # By setting scale_dim == None, weight_loader will
+ # repeat the scales on each rank in TP>1 case.
+ scales_and_zp_input_dim = None
+ scales_and_zp_size = input_size // group_size
+ else:
+ # By setting scale_dim == 0, weight_loader will
+ # shard the scales in TP>1 case.
+ scales_and_zp_input_dim = 0
+ scales_and_zp_size = input_size_per_partition // group_size
+
+ # Quantized weights
+ qweight = PackedvLLMParameter(
+ data=torch.empty(
+ input_size_per_partition // self.quant_config.pack_factor,
+ output_size_per_partition,
+ dtype=torch.int32,
+ ),
+ input_dim=0,
+ output_dim=1,
+ packed_dim=0,
+ packed_factor=self.quant_config.pack_factor,
+ weight_loader=weight_loader,
+ )
+
+ # Activation order
+ g_idx = RowvLLMParameter(
+ data=torch.empty(
+ input_size_per_partition,
+ dtype=torch.int32,
+ ),
+ input_dim=0,
+ weight_loader=weight_loader,
+ )
+ set_weight_attrs(
+ g_idx,
+ {"ignore_warning": True},
+ )
+
+ qzeros_args = {
+ "data": torch.empty(
+ scales_and_zp_size,
+ output_size_per_partition // self.quant_config.pack_factor,
+ dtype=torch.int32,
+ ),
+ "weight_loader": weight_loader,
+ }
+ weight_scale_args = {
+ "data": torch.empty(
+ scales_and_zp_size,
+ output_size_per_partition,
+ dtype=params_dtype,
+ ),
+ "weight_loader": weight_loader,
+ }
+
+ if scales_and_zp_input_dim is None:
+ scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
+ qzeros = PackedColumnParameter(
+ output_dim=1,
+ packed_dim=1,
+ packed_factor=self.quant_config.pack_factor,
+ **qzeros_args,
+ )
+
+ else:
+ scales = GroupQuantScaleParameter(
+ output_dim=1, input_dim=0, **weight_scale_args
+ )
+ qzeros = PackedvLLMParameter(
+ input_dim=0,
+ output_dim=1,
+ packed_dim=1,
+ packed_factor=self.quant_config.pack_factor,
+ **qzeros_args,
+ )
+
+ layer.register_parameter("qweight", qweight)
+ layer.register_parameter("g_idx", g_idx)
+ layer.register_parameter("scales", scales)
+ layer.register_parameter("qzeros", qzeros)
+
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+ torch.set_printoptions(profile="full", linewidth=5000, sci_mode=False)
+ packed_weight = layer.qweight.data
+ bits = self.quant_config.weight_bits
+ pack_factor = int(self.quant_config.pack_factor)
+ p_w_k, p_w_n = packed_weight.size()
+ input_size = p_w_k * pack_factor
+ output_size = p_w_n
+ isa_hint = _get_isa_hint(layer.scales.dtype)
+ layer.isa_hint = isa_hint
+
+ layer.qzeros = None
+ if not self.quant_config.desc_act:
+ layer.g_idx = None
+
+ # convert input dim packed to output dim packed
+ weight = unpack_cols(packed_weight, bits, p_w_k, p_w_n * pack_factor).view(
+ p_w_k, p_w_n, pack_factor
+ )
+ weight = weight.permute(0, 2, 1).reshape(input_size, output_size).contiguous()
+ weight = pack_cols(weight, bits, input_size, output_size)
+ # make 16 output channel as a block and transpose to the make
+ # the block contigous
+ weight = (
+ weight.view(input_size, -1, 16 // pack_factor)
+ .permute(1, 0, 2)
+ .reshape(-1, input_size * 16 // pack_factor)
+ .contiguous()
+ )
+ layer.qweight.data = weight
+
+ def apply(
+ self,
+ layer: torch.nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ x = cpu_gemm_wna16(
+ input=x,
+ q_weight=layer.qweight,
+ scales=layer.scales,
+ zeros=layer.qzeros,
+ g_idx=layer.g_idx,
+ bias=bias,
+ pack_factor=8,
+ isa_hint=layer.isa_hint,
+ )
+ return x
+
+
+class CPUAWQConfig(QuantizationConfig):
+ """Config class for CPU AWQ"""
+
+ def __init__(
+ self,
+ weight_bits: int,
+ group_size: int,
+ zero_point: bool,
+ lm_head_quantized: bool,
+ modules_to_not_convert: list[str] | None,
+ full_config: dict[str, Any],
+ ) -> None:
+ super().__init__()
+ assert weight_bits == 4
+ self.pack_factor = 32 // weight_bits # packed into int32
+ self.group_size = group_size
+ self.zero_point = zero_point
+ self.lm_head_quantized = lm_head_quantized
+ self.weight_bits = weight_bits
+ self.modules_to_not_convert = modules_to_not_convert or []
+ self.full_config = full_config
+
+ def __repr__(self) -> str:
+ return (
+ f"AWQMarlinConfig("
+ f"group_size={self.group_size}, "
+ f"zero_point={self.zero_point}, "
+ f"lm_head_quantized={self.lm_head_quantized}, "
+ f"modules_to_not_convert={self.modules_to_not_convert})"
+ )
+
+ @classmethod
+ def get_name(cls) -> "QuantizationMethods":
+ return "cpu_awq"
+
+ @classmethod
+ def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+ return [torch.half, torch.bfloat16]
+
+ @classmethod
+ def get_min_capability(cls) -> int:
+ return -1
+
+ @classmethod
+ def get_config_filenames(cls) -> list[str]:
+ return ["quantize_config.json"]
+
+ @classmethod
+ def from_config(cls, config: dict[str, Any]) -> "CPUAWQConfig":
+ weight_bits = cls.get_from_keys(config, ["bits"])
+ group_size = cls.get_from_keys(config, ["group_size"])
+ zero_point = cls.get_from_keys(config, ["zero_point"])
+ lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
+ modules_to_not_convert = cls.get_from_keys_or(
+ config, ["modules_to_not_convert"], None
+ )
+ return cls(
+ weight_bits,
+ group_size,
+ zero_point,
+ lm_head_quantized,
+ modules_to_not_convert,
+ config,
+ )
+
+ @classmethod
+ def override_quantization_method(
+ cls, hf_quant_cfg, user_quant
+ ) -> Optional["QuantizationMethods"]:
+ quant_method = hf_quant_cfg.get("quant_method", "").lower()
+ if current_platform.is_cpu() and (quant_method == "awq"):
+ return cls.get_name()
+ return None
+
+ def get_quant_method(
+ self, layer: torch.nn.Module, prefix: str
+ ) -> Optional["QuantizeMethodBase"]:
+ if isinstance(layer, LinearBase) or (
+ isinstance(layer, ParallelLMHead) and self.lm_head_quantized
+ ):
+ if is_layer_skipped(
+ prefix,
+ self.modules_to_not_convert,
+ self.packed_modules_mapping,
+ skip_with_substr=True,
+ ):
+ return UnquantizedLinearMethod()
+ return CPUAWQLinearMethod(self)
+ return None
+
+ def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
+ if self.modules_to_not_convert:
+ self.modules_to_not_convert = hf_to_vllm_mapper.apply_list(
+ self.modules_to_not_convert
+ )
+
+ def maybe_update_config(self, model_name: str, revision: str | None = None):
+ if self.modules_to_not_convert:
+ return
+
+ unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+ metadata = get_safetensors_params_metadata(model_name, revision=revision)
+ layers = {param_name.rsplit(".", 1)[0] for param_name in metadata}
+ quant_layers: set[str] = {
+ param_name.rsplit(".", 1)[0]
+ for param_name, info in metadata.items()
+ if (dtype := info.get("dtype", None))
+ and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
+ }
+ self.modules_to_not_convert = list(layers - quant_layers)
+
+
+class CPUAWQLinearMethod(LinearMethodBase):
+ """Linear method for CPU AWQ.
+
+ Args:
+ quant_config: The CPU AWQ quantization config.
+ """
+
+ def __init__(self, quant_config: CPUAWQConfig) -> None:
+ self.quant_config = quant_config
+ assert self.quant_config.zero_point
+
+ def create_weights(
+ self,
+ layer: torch.nn.Module,
+ input_size_per_partition: int,
+ output_partition_sizes: list[int],
+ input_size: int,
+ output_size: int,
+ params_dtype: torch.dtype,
+ **extra_weight_attrs,
+ ) -> None:
+ del output_size
+ output_size_per_partition = sum(output_partition_sizes)
+ weight_loader = extra_weight_attrs.get("weight_loader")
+
+ # Normalize group_size
+ if self.quant_config.group_size != -1:
+ group_size = self.quant_config.group_size
+ else:
+ group_size = input_size
+
+ qweight = PackedvLLMParameter(
+ data=torch.empty(
+ input_size_per_partition,
+ output_size_per_partition // self.quant_config.pack_factor,
+ dtype=torch.int32,
+ ),
+ input_dim=0,
+ output_dim=1,
+ packed_dim=1,
+ packed_factor=self.quant_config.pack_factor,
+ weight_loader=weight_loader,
+ )
+
+ num_groups = input_size_per_partition // group_size
+
+ qzeros = PackedvLLMParameter(
+ data=torch.empty(
+ num_groups,
+ output_size_per_partition // self.quant_config.pack_factor,
+ dtype=torch.int32,
+ ),
+ input_dim=0,
+ output_dim=1,
+ packed_dim=1,
+ packed_factor=self.quant_config.pack_factor,
+ weight_loader=weight_loader,
+ )
+
+ scales = GroupQuantScaleParameter(
+ data=torch.empty(
+ num_groups,
+ output_size_per_partition,
+ dtype=params_dtype,
+ ),
+ input_dim=0,
+ output_dim=1,
+ weight_loader=weight_loader,
+ )
+
+ layer.register_parameter("qweight", qweight)
+ layer.register_parameter("qzeros", qzeros)
+ layer.register_parameter("scales", scales)
+
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+ torch.set_printoptions(profile="full", linewidth=5000, sci_mode=False)
+ packed_weight = layer.qweight.data
+ packed_zeros = layer.qzeros.data
+ group_num = packed_zeros.size(0)
+ bits = self.quant_config.weight_bits
+ pack_factor = int(self.quant_config.pack_factor)
+ input_size, packed_output_size = packed_weight.size()
+ output_size = packed_output_size * pack_factor
+ isa_hint = _get_isa_hint(layer.scales.dtype)
+ layer.isa_hint = isa_hint
+
+ interleave_map = (0, 4, 1, 5, 2, 6, 3, 7)
+ weight = unpack_cols(
+ packed_weight,
+ bits,
+ input_size,
+ output_size,
+ )
+ zeros = unpack_cols(
+ packed_zeros,
+ bits,
+ group_num,
+ output_size,
+ )
+ weight = (
+ weight.view(input_size, -1, pack_factor)[:, :, interleave_map]
+ .reshape(input_size, output_size)
+ .contiguous()
+ )
+ zeros = (
+ zeros.view(group_num, -1, pack_factor)[:, :, interleave_map]
+ .reshape(group_num, output_size)
+ .contiguous()
+ )
+
+ zeros = pack_cols(zeros, bits, group_num, output_size).contiguous()
+ # make 16 output channel as a block and transpose to
+ # the make the block contigous
+ weight = pack_cols(weight, bits, input_size, output_size)
+ weight = (
+ weight.view(input_size, -1, 16 // pack_factor)
+ .permute(1, 0, 2)
+ .reshape(-1, input_size * 16 // pack_factor)
+ .contiguous()
+ )
+ layer.qweight.data = weight
+ layer.qzeros.data = zeros
+
+ def apply(
+ self,
+ layer: torch.nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ x = cpu_gemm_wna16(
+ input=x,
+ q_weight=layer.qweight,
+ scales=layer.scales,
+ zeros=layer.qzeros,
+ g_idx=None,
+ bias=bias,
+ pack_factor=8,
+ isa_hint=layer.isa_hint,
+ )
+ return x
+
+
+def _get_isa_hint(dtype: torch.dtype) -> str:
+ supports_amx = torch._C._cpu._is_amx_tile_supported()
+ if supports_amx and dtype in (torch.bfloat16,):
+ return "amx"
+ else:
+ return "vec"
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index e3a7b2a8dffe9..7fc074dfe478b 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -166,9 +166,11 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:
# deepGEMM on supported platforms with block-quantized weights
if envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM and block_quant:
if not has_deep_gemm():
- logger.warning_once("DeepGEMM backend requested but not available.")
+ logger.warning_once(
+ "DeepGEMM backend requested but not available.", scope="local"
+ )
elif is_deep_gemm_supported():
- logger.info_once("Using DeepGEMM backend for FP8 MoE")
+ logger.info_once("Using DeepGEMM backend for FP8 MoE", scope="local")
return Fp8MoeBackend.DEEPGEMM
# CUTLASS BlockScaled GroupedGemm on SM100 with block-quantized weights
@@ -177,7 +179,9 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:
and current_platform.is_device_capability(100)
and block_quant
):
- logger.info_once("Using Cutlass BlockScaled GroupedGemm backend for FP8 MoE")
+ logger.info_once(
+ "Using Cutlass BlockScaled GroupedGemm backend for FP8 MoE", scope="local"
+ )
return Fp8MoeBackend.CUTLASS_BLOCK_SCALED_GROUPED_GEMM
# default to Triton
@@ -1021,7 +1025,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
del layer.w13_input_scale
del layer.w2_input_scale
- def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:
+ def maybe_make_prepare_finalize(
+ self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+ ) -> mk.FusedMoEPrepareAndFinalize | None:
if (
self.rocm_aiter_moe_enabled
or self.use_marlin
@@ -1042,7 +1049,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
logger.debug_once("%s", prepare_finalize.__class__.__name__)
return prepare_finalize
else:
- return super().maybe_make_prepare_finalize()
+ return super().maybe_make_prepare_finalize(routing_tables)
def select_gemm_impl(
self,
diff --git a/vllm/model_executor/layers/quantization/gguf.py b/vllm/model_executor/layers/quantization/gguf.py
index caabcd0ca0ee5..42d7a67371ae8 100644
--- a/vllm/model_executor/layers/quantization/gguf.py
+++ b/vllm/model_executor/layers/quantization/gguf.py
@@ -1,7 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
+from types import MappingProxyType
from typing import Any, Optional
import gguf
@@ -26,7 +27,11 @@ from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+ UnquantizedEmbeddingMethod,
+ VocabParallelEmbedding,
+)
+from vllm.model_executor.models.utils import WeightsMapper
from vllm.model_executor.utils import set_weight_attrs
from vllm.utils.torch_utils import direct_register_custom_op
@@ -65,18 +70,70 @@ class GGUFConfig(QuantizationConfig):
self, layer: torch.nn.Module, prefix: str
) -> Optional["QuantizeMethodBase"]:
if isinstance(layer, LinearBase):
- if is_layer_skipped_gguf(prefix, self.unquantized_modules):
+ if is_layer_skipped_gguf(
+ prefix, self.unquantized_modules, self.packed_modules_mapping
+ ):
return UnquantizedLinearMethod()
return GGUFLinearMethod(self)
elif isinstance(layer, VocabParallelEmbedding):
+ if is_layer_skipped_gguf(
+ prefix, self.unquantized_modules, self.packed_modules_mapping
+ ):
+ return UnquantizedEmbeddingMethod()
return GGUFEmbeddingMethod(self)
elif isinstance(layer, FusedMoE):
return GGUFMoEMethod(self, layer.moe_config)
return None
+ def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
+ """
+ Interface for models to update module names referenced in
+ quantization configs in order to reflect the vllm model structure
-def is_layer_skipped_gguf(prefix: str, unquantized_modules: list[str]):
- return any(module_name in prefix for module_name in unquantized_modules)
+ :param hf_to_vllm_mapper: maps from hf model structure (the assumed
+ structure of the qconfig) to vllm model structure
+ """
+ if self.unquantized_modules is not None:
+ self.unquantized_modules = hf_to_vllm_mapper.apply_list(
+ self.unquantized_modules
+ )
+
+
+def is_layer_skipped_gguf(
+ prefix: str,
+ unquantized_modules: list[str],
+ fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
+):
+ # Fused layers like gate_up_proj or qkv_proj will not be fused
+ # in the safetensors checkpoint. So, we convert the name
+ # from the fused version to unfused + check to make sure that
+ # each shard of the fused layer has the same scheme.
+ proj_name = prefix.split(".")[-1]
+ if proj_name in fused_mapping:
+ shard_prefixes = [
+ prefix.replace(proj_name, shard_proj_name)
+ for shard_proj_name in fused_mapping[proj_name]
+ ]
+
+ is_skipped = None
+ for shard_prefix in shard_prefixes:
+ is_shard_skipped = any(
+ shard_prefix in module_name for module_name in unquantized_modules
+ )
+
+ if is_skipped is None:
+ is_skipped = is_shard_skipped
+ elif is_shard_skipped != is_skipped:
+ raise ValueError(
+ f"Detected some but not all shards of {prefix} "
+ "are quantized. All shards of fused layers "
+ "to have the same precision."
+ )
+ else:
+ is_skipped = any(module_name in prefix for module_name in unquantized_modules)
+
+ assert is_skipped is not None
+ return is_skipped
UNQUANTIZED_TYPES = {WeightType.F32, WeightType.F16, WeightType.BF16}
diff --git a/vllm/model_executor/layers/quantization/ipex_quant.py b/vllm/model_executor/layers/quantization/ipex_quant.py
index 5ca9167faec80..22c4bae041a56 100644
--- a/vllm/model_executor/layers/quantization/ipex_quant.py
+++ b/vllm/model_executor/layers/quantization/ipex_quant.py
@@ -134,7 +134,7 @@ class IPEXConfig(QuantizationConfig):
def override_quantization_method(
cls, hf_quant_cfg, user_quant
) -> QuantizationMethods | None:
- if not current_platform.is_cpu() and not current_platform.is_xpu():
+ if not current_platform.is_xpu():
return None
quant_method = hf_quant_cfg.get("quant_method", "").lower()
diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py
index 9309f4f150e45..d1f8f39a69b3b 100644
--- a/vllm/model_executor/layers/quantization/modelopt.py
+++ b/vllm/model_executor/layers/quantization/modelopt.py
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
+from fnmatch import fnmatch
from typing import TYPE_CHECKING, Any, Optional
import torch
@@ -13,9 +14,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.config import (
- FusedMoEConfig,
FusedMoEQuantConfig,
- RoutingMethodType,
fp8_w8a8_moe_quant_config,
nvfp4_moe_quant_config,
)
@@ -41,6 +40,8 @@ from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
build_flashinfer_fp4_cutlass_moe_prepare_finalize,
+ flashinfer_trtllm_fp4_moe,
+ prepare_static_weights_for_trtllm_fp4_moe,
reorder_w1w3_to_w3w1,
select_nvfp4_gemm_impl,
)
@@ -88,45 +89,218 @@ QUANT_ALGOS = ["FP8", "NVFP4"]
KV_CACHE_QUANT_ALGOS = ["FP8"]
-class ModelOptFp8Config(QuantizationConfig):
+class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
+ """
+ Supports loading kv-cache scaling factors from FP8 checkpoints.
+ """
+
+ def __init__(self, quant_config: "ModelOptQuantConfigBase"):
+ super().__init__(quant_config)
+
+
+class ModelOptQuantConfigBase(QuantizationConfig):
+ LinearMethodCls: type = LinearMethodBase
+ FusedMoEMethodCls: type = FusedMoEMethodBase
+ KVCacheMethodCls: type = BaseKVCacheMethod
+
+ def __init__(
+ self,
+ exclude_modules: list[str],
+ ):
+ super().__init__()
+ self.exclude_modules: list[str] = exclude_modules
+
+ def is_layer_excluded(self, prefix: str) -> bool:
+ """
+ Check if a layer should be excluded from quantization.
+
+ Handles both exact matching (for fused layers) and ModelOpt wildcard matching.
+
+ The ModelOpt exclude_modules list is a list of wildcards.
+ """
+ if len(self.exclude_modules) == 0:
+ return False
+
+ # First check exact matching with fused layer support
+ if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):
+ return True
+
+ # TODO: This special hard coded logic is not needed for quantized checkpoints
+ # generated by ModelOpt >= 0.39.0 where they are handled natually by the
+ # exclude_modules config. But need to keep them for loading quantized
+ # checkpoints generated by older versions. Then check substring matching
+ # for patterns not caught by exact match
+ for exclude_module in self.exclude_modules:
+ # Skip exact matches already handled above
+ if exclude_module != prefix and (
+ exclude_module in prefix
+ or (
+ prefix.startswith("language_model.")
+ and exclude_module in prefix.removeprefix("language_model.")
+ )
+ ):
+ return True
+
+ # modelopt exclude modules are not simple strings, they are wildcards
+ for wildcard_pattern in self.exclude_modules:
+ if fnmatch(prefix, wildcard_pattern):
+ return True
+
+ return False
+
+ def get_quant_method(
+ self, layer: torch.nn.Module, prefix: str
+ ) -> Optional["QuantizeMethodBase"]:
+ from vllm.attention.layer import Attention # Avoid circular import
+
+ # handle kv-cache first so we can focus only on weight quantization thereafter
+ if isinstance(layer, Attention):
+ return self.KVCacheMethodCls(self)
+
+ # handle exclusion
+ if self.is_layer_excluded(prefix):
+ if isinstance(layer, LinearBase):
+ return UnquantizedLinearMethod()
+ return None
+
+ # TODO: This special hard coded logic is not needed for quantized checkpoints
+ # generated by ModelOpt >= 0.39.0 where they are handled natually by the
+ # exclude_modules config. But need to keep them for loading quantized
+ # checkpoints generated by older versions. Then check substring matching
+ # for patterns not caught by exact match
+ if "vision_tower" in prefix or "vision_model" in prefix:
+ return UnquantizedLinearMethod()
+
+ # now, the layer is quantized, handle it here
+ if isinstance(layer, LinearBase):
+ return self.LinearMethodCls(self)
+ elif isinstance(layer, FusedMoE):
+ return self.FusedMoEMethodCls(quant_config=self, layer=layer)
+
+ return None
+
+ def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
+ if len(self.exclude_modules) > 0:
+ self.exclude_modules = hf_to_vllm_mapper.apply_list(self.exclude_modules)
+
+ @staticmethod
+ def get_config_filenames() -> list[str]:
+ return ["hf_quant_config.json"]
+
+ @classmethod
+ def _from_config(
+ cls,
+ *,
+ quant_method: str,
+ kv_cache_quant_method: str | None,
+ exclude_modules: list[str],
+ original_config: dict[str, Any],
+ group_size: int | None,
+ ) -> "ModelOptQuantConfigBase":
+ raise NotImplementedError("Please implement this function in sub classes")
+
+ @classmethod
+ def from_config(cls, config: dict[str, Any]) -> "ModelOptQuantConfigBase":
+ # Handle both ModelOpt format and compressed-tensors style format
+ if "quantization" in config:
+ # Traditional ModelOpt format:
+ # {"quantization": {"quant_algo": "..."}}
+ quant_config = cls.get_from_keys(config, ["quantization"])
+ if not isinstance(quant_config, dict):
+ raise ValueError("Expected 'quantization' to be a dictionary in config")
+
+ quant_method = quant_config.get("quant_algo")
+
+ # Handle kv_cache_quant_algo with proper type validation
+ kv_cache_quant_method = quant_config.get("kv_cache_quant_algo")
+
+ # Handle group_size with proper type validation
+ group_size_raw = quant_config.get("group_size")
+
+ # "exclude_modules" is the key in the legacy hf_quant_config.json
+ exclude_modules = quant_config.get("exclude_modules", [])
+ else:
+ # Compressed-tensors style format:
+ # {"quant_algo": "...", "quant_method": "modelopt"}
+ quant_method = config.get("quant_algo")
+ kv_cache_quant_method = config.get("kv_cache_quant_algo")
+ # "ignore" is the key in config.json
+ exclude_modules = config.get("ignore", [])
+ group_size_raw = config.get("group_size")
+
+ if not quant_method:
+ raise ValueError("Missing 'quant_algo' in quantization config")
+
+ if kv_cache_quant_method is None:
+ # No KV cache quantization, keep this branch just to have this comment
+ pass
+ elif not isinstance(kv_cache_quant_method, str):
+ raise ValueError(
+ f"kv_cache_quant_algo must be a string, got "
+ f"{type(kv_cache_quant_method)}"
+ )
+
+ if not isinstance(exclude_modules, list):
+ raise ValueError(
+ f"exclude_modules must be a list, got {type(exclude_modules)}"
+ )
+
+ if group_size_raw is None:
+ group_size = None
+ elif isinstance(group_size_raw, int):
+ group_size = group_size_raw
+ else:
+ try:
+ group_size = int(group_size_raw)
+ except (ValueError, TypeError):
+ raise ValueError(
+ f"group_size must be an integer, got {type(group_size_raw)}"
+ ) from None
+
+ if quant_method not in QUANT_ALGOS:
+ raise ValueError(
+ f"ModelOpt currently only supports: {QUANT_ALGOS} "
+ "quantizations in vLLM. Please check the "
+ "`hf_quant_config.json` file for your model's "
+ "quant configuration."
+ )
+ return cls._from_config(
+ quant_method=quant_method,
+ kv_cache_quant_method=kv_cache_quant_method,
+ exclude_modules=exclude_modules,
+ group_size=group_size,
+ original_config=config,
+ )
+
+
+class ModelOptFp8Config(ModelOptQuantConfigBase):
"""Config class for ModelOpt FP8."""
def __init__(
self,
- is_checkpoint_fp8_serialized: bool = False,
- kv_cache_quant_method: str | None = None,
- exclude_modules: list[str] | None = None,
+ is_checkpoint_fp8_serialized: bool,
+ kv_cache_quant_method: str | None,
+ exclude_modules: list[str],
) -> None:
- super().__init__()
+ super().__init__(exclude_modules)
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
self.kv_cache_quant_method = kv_cache_quant_method
- self.exclude_modules = exclude_modules or []
if is_checkpoint_fp8_serialized:
logger.warning(
"Detected ModelOpt fp8 checkpoint. Please note that"
" the format is experimental and could change."
)
- @classmethod
- def get_name(cls) -> QuantizationMethods:
+ def get_name(self) -> QuantizationMethods:
return "modelopt"
- @classmethod
- def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+ def get_supported_act_dtypes(self) -> list[torch.dtype]:
return [torch.bfloat16, torch.half]
@classmethod
def get_min_capability(cls) -> int:
return 89
- @classmethod
- def get_config_filenames(cls) -> list[str]:
- return ["hf_quant_config.json"]
-
- def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
- if self.exclude_modules is not None:
- self.exclude_modules = hf_to_vllm_mapper.apply_list(self.exclude_modules)
-
@classmethod
def override_quantization_method(
cls, hf_quant_cfg, user_quant
@@ -160,88 +334,19 @@ class ModelOptFp8Config(QuantizationConfig):
return None
@classmethod
- def from_config(cls, config: dict[str, Any]) -> "ModelOptFp8Config":
- # Handle both ModelOpt format and compressed-tensors style format
- if "quantization" in config:
- # ModelOpt format: {"quantization": {"quant_algo": "..."}}
- quant_config = cls.get_from_keys(config, ["quantization"])
- if not isinstance(quant_config, dict):
- raise ValueError("Expected 'quantization' to be a dictionary in config")
- quant_method = quant_config.get("quant_algo", "")
- if not quant_method:
- raise ValueError("Missing 'quant_algo' in quantization config")
- kv_cache_quant_method = quant_config.get("kv_cache_quant_algo")
- # "exclude_modules" is the key in the legacy hf_quant_config.json
- exclude_modules = quant_config.get("exclude_modules")
- else:
- # Compressed-tensors style format:
- # {"quant_algo": "...", "quant_method": "modelopt"}
- quant_method = config.get("quant_algo", "")
- kv_cache_quant_method = config.get("kv_cache_quant_algo")
- # "ignore" is the key in config.json
- exclude_modules = config.get("ignore")
-
- if quant_method not in QUANT_ALGOS:
- raise ValueError(
- f"ModelOpt currently only supports: {QUANT_ALGOS} "
- "quantizations in vLLM. Please check the "
- "`hf_quant_config.json` file for your model's "
- "quant configuration."
- )
+ def _from_config(
+ cls,
+ *,
+ quant_method: str,
+ kv_cache_quant_method: str | None,
+ exclude_modules: list[str],
+ original_config: dict[str, Any],
+ **kwargs: Any,
+ ) -> "ModelOptFp8Config":
is_checkpoint_fp8_serialized = "FP8" in quant_method
return cls(is_checkpoint_fp8_serialized, kv_cache_quant_method, exclude_modules)
- def is_layer_excluded(self, prefix: str) -> bool:
- """
- Check if a layer should be excluded from quantization.
- Handles both exact matching (for fused layers) and substring matching.
-
- This method handles both regular models and multimodal models that use
- the language_model prefix. For multimodal models, it checks if the
- module name (without the language_model prefix) is in the exclude list.
- """
- if self.exclude_modules is None:
- return False
-
- # First check exact matching with fused layer support
- if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):
- return True
-
- # Then check substring matching for patterns not caught by exact match
- for module in self.exclude_modules:
- # Skip exact matches already handled above
- if module != prefix and (
- module in prefix
- or (
- prefix.startswith("language_model.")
- and module in prefix.removeprefix("language_model.")
- )
- ):
- return True
- return False
-
- def get_quant_method(
- self, layer: torch.nn.Module, prefix: str
- ) -> Optional["QuantizeMethodBase"]:
- from vllm.attention.layer import ( # Avoid circular import
- Attention,
- MLAAttention,
- )
-
- if isinstance(layer, LinearBase):
- if self.is_layer_excluded(prefix):
- return UnquantizedLinearMethod()
- # Check if this is a vision model layer that should not be quantized
- if "vision_tower" in prefix or "vision_model" in prefix:
- return UnquantizedLinearMethod()
- return ModelOptFp8LinearMethod(self)
- elif isinstance(layer, (Attention, MLAAttention)):
- return ModelOptFp8KVCacheMethod(self)
- elif isinstance(layer, FusedMoE):
- return ModelOptFp8MoEMethod(self, layer)
- return None
-
class ModelOptFp8LinearMethod(LinearMethodBase):
"""Linear method for Model Optimizer static quantization.
@@ -343,7 +448,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
def __init__(
self,
quant_config: ModelOptFp8Config,
- layer: torch.nn.Module,
+ layer: FusedMoE,
) -> None:
super().__init__(layer.moe_config)
self.layer = layer
@@ -372,6 +477,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
def maybe_make_prepare_finalize(
self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> mk.FusedMoEPrepareAndFinalize | None:
# TRT LLM not supported with all2all yet.
if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
@@ -383,7 +489,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
logger.debug_once("%s", prepare_finalize.__class__.__name__)
return prepare_finalize
else:
- return super().maybe_make_prepare_finalize()
+ return super().maybe_make_prepare_finalize(routing_tables)
def select_gemm_impl(
self,
@@ -684,7 +790,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
)
-class ModelOptNvFp4Config(QuantizationConfig):
+ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod
+ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod
+ModelOptFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
+
+
+class ModelOptNvFp4Config(ModelOptQuantConfigBase):
"""Config class for ModelOpt FP4."""
def __init__(
@@ -694,7 +805,7 @@ class ModelOptNvFp4Config(QuantizationConfig):
exclude_modules: list[str],
group_size: int = 16,
) -> None:
- super().__init__()
+ super().__init__(exclude_modules)
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
if is_checkpoint_nvfp4_serialized:
logger.warning(
@@ -704,28 +815,17 @@ class ModelOptNvFp4Config(QuantizationConfig):
self.group_size = group_size
self.kv_cache_quant_algo = kv_cache_quant_algo
- self.exclude_modules = exclude_modules
- @classmethod
- def get_name(cls) -> QuantizationMethods:
+ def get_name(self) -> QuantizationMethods:
return "modelopt_fp4"
- @classmethod
- def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+ def get_supported_act_dtypes(self) -> list[torch.dtype]:
return [torch.bfloat16, torch.half, torch.float8_e4m3fn]
@classmethod
def get_min_capability(cls) -> int:
return 80
- @classmethod
- def get_config_filenames(cls) -> list[str]:
- return ["hf_quant_config.json"]
-
- def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
- if self.exclude_modules is not None:
- self.exclude_modules = hf_to_vllm_mapper.apply_list(self.exclude_modules)
-
@classmethod
def override_quantization_method(
cls, hf_quant_cfg, user_quant
@@ -759,105 +859,25 @@ class ModelOptNvFp4Config(QuantizationConfig):
return None
@classmethod
- def from_config(cls, config: dict[str, Any]) -> "ModelOptNvFp4Config":
- # Handle both traditional ModelOpt format and compressed-tensors
- # style format
- if "quantization" in config:
- # Traditional ModelOpt format:
- # {"quantization": {"quant_algo": "..."}}
- quant_config = cls.get_from_keys(config, ["quantization"])
- if not isinstance(quant_config, dict):
- raise ValueError("Expected 'quantization' to be a dictionary in config")
-
- quant_method = quant_config.get("quant_algo", "")
- if not quant_method:
- raise ValueError("Missing 'quant_algo' in quantization config")
-
- # Handle kv_cache_quant_algo with proper type validation
- kv_cache_quant_algo_raw = quant_config.get("kv_cache_quant_algo")
- if kv_cache_quant_algo_raw is None:
- # No KV cache quantization by default
- kv_cache_quant_algo = None
- elif isinstance(kv_cache_quant_algo_raw, str):
- kv_cache_quant_algo = kv_cache_quant_algo_raw
- else:
- raise ValueError(
- f"kv_cache_quant_algo must be a string, got "
- f"{type(kv_cache_quant_algo_raw)}"
- )
-
- # Handle group_size with proper type validation
- group_size_raw = quant_config.get("group_size")
- if group_size_raw is None:
- group_size = 16 # Default value
- elif isinstance(group_size_raw, int):
- group_size = group_size_raw
- else:
- try:
- group_size = int(group_size_raw)
- except (ValueError, TypeError):
- raise ValueError(
- f"group_size must be an integer, got {type(group_size_raw)}"
- ) from None
-
- # "exclude_modules" is the key in the legacy hf_quant_config.json
- exclude_modules = quant_config.get("exclude_modules", [])
- if not isinstance(exclude_modules, list):
- raise ValueError(
- f"exclude_modules must be a list, got {type(exclude_modules)}"
- )
- else:
- # Compressed-tensors style format:
- # {"quant_algo": "...", "quant_method": "modelopt"}
- quant_method = config.get("quant_algo", "")
-
- # Handle kv_cache_quant_algo with proper type validation
- kv_cache_quant_algo_raw = config.get("kv_cache_quant_algo")
- if kv_cache_quant_algo_raw is None:
- # No KV cache quantization by default
- kv_cache_quant_algo = None
- elif isinstance(kv_cache_quant_algo_raw, str):
- kv_cache_quant_algo = kv_cache_quant_algo_raw
- else:
- raise ValueError(
- f"kv_cache_quant_algo must be a string, got "
- f"{type(kv_cache_quant_algo_raw)}"
- )
-
- # Handle group_size with proper type validation
- group_size_raw = config.get("group_size")
- if group_size_raw is None:
- group_size = 16 # Default value
- elif isinstance(group_size_raw, int):
- group_size = group_size_raw
- else:
- try:
- group_size = int(group_size_raw)
- except (ValueError, TypeError):
- raise ValueError(
- f"group_size must be an integer, got {type(group_size_raw)}"
- ) from None
-
- # "ignore" is the key in config.json
- exclude_modules = config.get("ignore", [])
- if not isinstance(exclude_modules, list):
- raise ValueError(
- f"exclude_modules must be a list, got {type(exclude_modules)}"
- )
-
- if quant_method not in QUANT_ALGOS:
- raise ValueError(
- f"ModelOpt currently only supports: {QUANT_ALGOS} "
- "quantizations in vLLM. Please check the "
- "`hf_quant_config.json` file for your model's "
- "quant configuration."
- )
+ def _from_config(
+ cls,
+ *,
+ quant_method: str,
+ kv_cache_quant_method: str | None,
+ exclude_modules: list[str],
+ original_config: dict[str, Any],
+ group_size: int | None,
+ **kwargs: Any,
+ ) -> "ModelOptNvFp4Config":
is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method
+ if group_size is None:
+ group_size = 16 # Default value
+
# For FP4, these fields are required
- if is_checkpoint_nvfp4_serialized and "quantization" in config:
+ if is_checkpoint_nvfp4_serialized and "quantization" in original_config:
# Check if required fields are present in the quantization config
- quant_config = config["quantization"]
+ quant_config = original_config["quantization"]
required_fields = ["group_size", "kv_cache_quant_algo", "exclude_modules"]
missing_fields = [
field for field in required_fields if field not in quant_config
@@ -870,64 +890,11 @@ class ModelOptNvFp4Config(QuantizationConfig):
return cls(
is_checkpoint_nvfp4_serialized,
- kv_cache_quant_algo,
+ kv_cache_quant_method,
exclude_modules,
group_size,
)
- def is_layer_excluded(self, prefix: str) -> bool:
- """
- Check if a layer should be excluded from quantization.
- Handles both exact matching (for fused layers) and pattern matching.
- """
- # First check exact matching with fused layer support
- if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):
- return True
-
- # Check regex pattern matching for patterns not caught by exact match
- import regex as re
-
- for pattern in self.exclude_modules:
- # Skip patterns that would be caught by exact matching
- if "*" in pattern or "." in pattern:
- regex_str = pattern.replace(".", r"\.").replace("*", r".*")
- if re.fullmatch(regex_str, prefix):
- return True
- return False
-
- def get_quant_method(
- self, layer: torch.nn.Module, prefix: str
- ) -> Optional["QuantizeMethodBase"]:
- from vllm.attention.layer import ( # Avoid circular import
- Attention,
- MLAAttention,
- )
-
- skip_layer = self.is_layer_excluded(prefix)
- if isinstance(layer, LinearBase):
- if skip_layer:
- return UnquantizedLinearMethod()
- # Check if this is a vision model layer that should not be quantized
- if "vision_tower" in prefix or "vision_model" in prefix:
- return UnquantizedLinearMethod()
- return ModelOptNvFp4LinearMethod(self)
- elif isinstance(layer, (Attention, MLAAttention)):
- return ModelOptFp8KVCacheMethod(self)
- elif isinstance(layer, FusedMoE):
- if skip_layer:
- return None
- return ModelOptNvFp4FusedMoE(self, layer.moe_config, layer)
- return None
-
-
-class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
- """
- Supports loading kv-cache scaling factors from FP8 checkpoints.
- """
-
- def __init__(self, quant_config: ModelOptFp8Config | ModelOptNvFp4Config):
- super().__init__(quant_config)
-
class ModelOptNvFp4LinearMethod(LinearMethodBase):
"""Linear method for Model Optimizer NVFP4.
@@ -1155,14 +1122,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
def __init__(
self,
quant_config: ModelOptNvFp4Config,
- moe: FusedMoEConfig,
- layer: torch.nn.Module,
+ layer: FusedMoE,
) -> None:
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import (
detect_nvfp4_moe_support, # noqa: E501
)
- super().__init__(moe)
+ super().__init__(layer.moe_config)
self.quant_config = quant_config
self.layer = layer
_nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)
@@ -1170,7 +1136,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
self.allow_flashinfer = _nvfp4.allow_flashinfer
self.use_marlin = _nvfp4.use_marlin
self.flashinfer_moe_backend = None
- self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
if self.allow_flashinfer:
self.flashinfer_moe_backend = get_flashinfer_moe_backend()
logger.info_once(
@@ -1178,7 +1143,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
" for ModelOptNvFp4FusedMoE."
)
- def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:
+ def maybe_make_prepare_finalize(
+ self,
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+ ) -> mk.FusedMoEPrepareAndFinalize | None:
if self.use_marlin or (
self.allow_flashinfer
and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
@@ -1195,7 +1163,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
logger.debug_once("%s", prepare_finalize.__class__.__name__)
return prepare_finalize
else:
- return super().maybe_make_prepare_finalize()
+ return super().maybe_make_prepare_finalize(routing_tables)
def select_gemm_impl(
self,
@@ -1334,136 +1302,15 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
)
layer.register_parameter("w2_input_scale", w2_input_scale)
- def prepare_static_weights_for_trtllm_fp4_moe(
- self,
- # args_dequant,
- # args,
- gemm1_weights,
- gemm2_weights,
- gemm1_scales_linear_fp4_bytes,
- gemm2_scales_linear_fp4_bytes,
- hidden_size,
- intermediate_size,
- num_experts,
- ):
- from flashinfer import nvfp4_block_scale_interleave
- from flashinfer.fused_moe.core import (
- _maybe_get_cached_w3_w1_permute_indices,
- get_w2_permute_indices_with_cache,
- )
-
- """Prepare quantized weights for kernel (done offline with weights)."""
- epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
-
- # Convert quantized weights to proper formats
- gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
- num_experts, 2 * intermediate_size, hidden_size // 2
- ) # packed fp4
- gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
- torch.float8_e4m3fn
- ).reshape(
- num_experts, 2 * intermediate_size, hidden_size // 16
- ) # fp8 scaling factors
-
- gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
- num_experts, hidden_size, intermediate_size // 2
- ) # packed fp4
- gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(
- torch.float8_e4m3fn
- ).reshape(
- num_experts, hidden_size, intermediate_size // 16
- ) # fp8 scaling factors
-
- gemm1_weights_fp4_shuffled = []
- gemm1_scales_fp4_shuffled = []
- gemm2_weights_fp4_shuffled = []
- gemm2_scales_fp4_shuffled = []
- for i in range(num_experts):
- # Calculate the permute indices for the following:
- # 1. Reorder rows of W1 and scales for fused gated activation
- # 2. Shuffle weights and scaling factors for transposed mma output
- # for both w3_w1 and w2 weights and scale factors
- permute_indices = _maybe_get_cached_w3_w1_permute_indices(
- self._cache_permute_indices,
- gemm1_weights_fp4[i].view(torch.uint8),
- epilogue_tile_m,
- )
- gemm1_weights_fp4_shuffled.append(
- gemm1_weights_fp4[i]
- .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]
- .contiguous()
- )
-
- permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
- self._cache_permute_indices,
- gemm1_scales_linear_fp4[i].view(torch.uint8),
- epilogue_tile_m,
- num_elts_per_sf=16,
- )
- gemm1_scales_fp4_shuffled.append(
- nvfp4_block_scale_interleave(
- gemm1_scales_linear_fp4[i]
- .view(torch.uint8)[
- permute_sf_indices.to(gemm1_scales_linear_fp4.device)
- ]
- .contiguous()
- )
- )
-
- permute_indices = get_w2_permute_indices_with_cache(
- self._cache_permute_indices,
- gemm2_weights_fp4[i].view(torch.uint8),
- epilogue_tile_m,
- )
- gemm2_weights_fp4_shuffled.append(
- gemm2_weights_fp4[i]
- .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]
- .contiguous()
- )
-
- permute_sf_indices = get_w2_permute_indices_with_cache(
- self._cache_permute_indices,
- gemm2_scales_linear_fp4[i].view(torch.uint8),
- epilogue_tile_m,
- num_elts_per_sf=16,
- )
- gemm2_scales_fp4_shuffled.append(
- nvfp4_block_scale_interleave(
- gemm2_scales_linear_fp4[i]
- .view(torch.uint8)[
- permute_sf_indices.to(gemm2_scales_linear_fp4.device)
- ]
- .contiguous()
- )
- )
-
- # Stack weights for all experts
- gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
- gemm1_scales_fp4_shuffled = (
- torch.stack(gemm1_scales_fp4_shuffled)
- .view(torch.float8_e4m3fn)
- .reshape(num_experts, 2 * intermediate_size, hidden_size // 16)
- )
-
- gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)
- gemm2_scales_fp4_shuffled = (
- torch.stack(gemm2_scales_fp4_shuffled)
- .view(torch.float8_e4m3fn)
- .reshape(num_experts, hidden_size, intermediate_size // 16)
- )
- return (
- gemm1_weights_fp4_shuffled,
- gemm1_scales_fp4_shuffled,
- gemm2_weights_fp4_shuffled,
- gemm2_scales_fp4_shuffled,
- )
-
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# GEMM 1 processing
gemm1_weight = layer.w13_weight.data
gemm1_weight_scale = layer.w13_weight_scale.data
- if self.allow_flashinfer:
+ if self.allow_flashinfer and (
+ self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
+ or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+ ):
gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(
gemm1_weight, gemm1_weight_scale, dim=-2
)
@@ -1536,7 +1383,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
gemm1_scales_fp4_shuffled,
gemm2_weights_fp4_shuffled,
gemm2_scales_fp4_shuffled,
- ) = self.prepare_static_weights_for_trtllm_fp4_moe(
+ ) = prepare_static_weights_for_trtllm_fp4_moe(
layer.w13_weight,
layer.w2_weight,
layer.w13_weight_scale,
@@ -1642,68 +1489,17 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
self.allow_flashinfer
and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
):
- import flashinfer
-
- from vllm.model_executor.models.llama4 import Llama4MoE
-
- a1_gscale = layer.w13_input_scale_quant
- (hidden_states_fp4, hidden_states_scale_linear_fp4) = (
- flashinfer.fp4_quantize(
- x,
- a1_gscale,
- is_sf_swizzled_layout=False,
- )
- )
- use_llama4_routing = (
- custom_routing_function is Llama4MoE.custom_routing_function
- )
- routing_method_type = layer.routing_method_type
- if use_llama4_routing:
- routing_method_type = RoutingMethodType.Llama4
- router_logits = (
- router_logits.to(torch.float32)
- if routing_method_type == RoutingMethodType.DeepSeekV3
- else router_logits
- )
- routing_bias = e_score_correction_bias
- if routing_bias is not None:
- routing_bias = routing_bias.to(torch.bfloat16)
- out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
- routing_logits=router_logits,
- routing_bias=routing_bias,
- hidden_states=hidden_states_fp4,
- hidden_states_scale=hidden_states_scale_linear_fp4.view(
- torch.float8_e4m3fn
- ).flatten(),
- gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,
- gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(
- torch.float8_e4m3fn
- ),
- gemm1_bias=None,
- gemm1_alpha=None,
- gemm1_beta=None,
- gemm1_clamp_limit=None,
- gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,
- gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(
- torch.float8_e4m3fn
- ),
- gemm2_bias=None,
- output1_scale_scalar=layer.g1_scale_c.data,
- output1_scale_gate_scalar=layer.g1_alphas.data,
- output2_scale_scalar=layer.g2_alphas.data,
- num_experts=global_num_experts,
+ return flashinfer_trtllm_fp4_moe(
+ layer=layer,
+ x=x,
+ router_logits=router_logits,
top_k=top_k,
- n_group=num_expert_group,
+ global_num_experts=global_num_experts,
+ num_expert_group=num_expert_group,
topk_group=topk_group,
- intermediate_size=layer.intermediate_size_per_partition,
- local_expert_offset=layer.ep_rank * layer.local_num_experts,
- local_num_experts=layer.local_num_experts,
- routed_scaling_factor=None,
- tile_tokens_dim=None,
- routing_method_type=routing_method_type,
- do_finalize=True,
- )[0]
- return out
+ custom_routing_function=custom_routing_function,
+ e_score_correction_bias=e_score_correction_bias,
+ )
topk_weights, topk_ids, _ = FusedMoE.select_experts(
hidden_states=x,
@@ -1741,17 +1537,26 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
workspace=layer.workspace,
)
- elif (
- self.allow_flashinfer
- and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
- ):
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
- flashinfer_cutlass_moe_fp4,
+ elif self.allow_flashinfer:
+ assert self.flashinfer_moe_backend in (
+ FlashinferMoeBackend.CUTLASS,
+ FlashinferMoeBackend.CUTEDSL,
)
+ if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
+ flashinfer_cutlass_moe_fp4,
+ )
+
+ flashinfer_fn_moe_fp4 = flashinfer_cutlass_moe_fp4
+ else:
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import ( # noqa: E501
+ flashinfer_cutedsl_moe_fp4,
+ )
+
+ flashinfer_fn_moe_fp4 = flashinfer_cutedsl_moe_fp4
assert self.moe_quant_config is not None
-
- return flashinfer_cutlass_moe_fp4(
+ return flashinfer_fn_moe_fp4(
hidden_states=x,
w1=layer.w13_weight,
w2=layer.w2_weight,
@@ -1785,3 +1590,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
k=x.shape[1],
e=layer.w13_weight.shape[0],
)
+
+
+ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod
+ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE
+ModelOptNvFp4Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py
index b95d1a6b3a1f5..66ae2e94c60a5 100644
--- a/vllm/model_executor/layers/quantization/mxfp4.py
+++ b/vllm/model_executor/layers/quantization/mxfp4.py
@@ -755,8 +755,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self.w13_weight = w13_weight
self.w2_weight = w2_weight
- layer.w13_weight = Parameter(w13_weight.storage.data, requires_grad=False)
- layer.w2_weight = Parameter(w2_weight.storage.data, requires_grad=False)
+ del layer.w13_weight
+ del layer.w2_weight
+ layer.w13_weight = w13_weight
+ layer.w2_weight = w2_weight
else:
raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
@@ -1065,8 +1067,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
return triton_kernel_moe_forward(
hidden_states=x,
- w1=self.w13_weight,
- w2=self.w2_weight,
+ w1=layer.w13_weight,
+ w2=layer.w2_weight,
gating_output=router_logits,
topk=top_k,
renormalize=renormalize,
diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py
index 1bb698faf46df..f59e5e2a0af7a 100644
--- a/vllm/model_executor/layers/quantization/quark/quark.py
+++ b/vllm/model_executor/layers/quantization/quark/quark.py
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.quark.utils import (
deep_compare,
should_ignore_layer,
)
+from vllm.model_executor.models.utils import WeightsMapper
from vllm.platforms import current_platform
if TYPE_CHECKING:
@@ -57,7 +58,6 @@ class QuarkConfig(QuantizationConfig):
self.kv_cache_group = kv_cache_group
self.kv_cache_config = kv_cache_config
self.pack_method = pack_method
- self.ignore: list[str] = cast(list[str], self.quant_config.get("exclude", []))
def get_linear_method(self) -> "QuarkLinearMethod":
return QuarkLinearMethod(self)
@@ -72,14 +72,42 @@ class QuarkConfig(QuantizationConfig):
def get_name(self) -> QuantizationMethods:
return "quark"
+ def apply_vllm_mapper( # noqa: B027
+ self, hf_to_vllm_mapper: "WeightsMapper"
+ ):
+ """
+ Interface for models to update module names referenced in
+ quantization configs in order to reflect the vllm model structure
+
+ :param hf_to_vllm_mapper: maps from hf model structure (the assumed
+ structure of the qconfig) to vllm model structure
+ """
+ quant_config_with_hf_to_vllm_mapper = {}
+
+ for k, v in self.quant_config.items():
+ if isinstance(v, list):
+ quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_list(v)
+ elif isinstance(v, dict):
+ quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_dict(v)
+ else:
+ if isinstance(v, str):
+ mapped_v_list = hf_to_vllm_mapper.apply_list([v])
+ if mapped_v_list:
+ quant_config_with_hf_to_vllm_mapper[k] = mapped_v_list[0]
+ else:
+ quant_config_with_hf_to_vllm_mapper[k] = v
+
+ self.quant_config = quant_config_with_hf_to_vllm_mapper
+
def get_quant_method(
self, layer: torch.nn.Module, prefix: str
) -> Optional["QuantizeMethodBase"]:
from vllm.attention.layer import Attention # Avoid circular import
# Check if the layer is skipped for quantization.
+ exclude_layers = cast(list[str], self.quant_config.get("exclude"))
if should_ignore_layer(
- prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping
+ prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping
):
return UnquantizedLinearMethod()
if isinstance(layer, LinearBase):
@@ -93,9 +121,6 @@ class QuarkConfig(QuantizationConfig):
return QuarkMoEMethod.get_moe_method(self, module=layer, layer_name=prefix)
return None
- def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
- self.ignore = hf_to_vllm_mapper.apply_list(self.ignore)
-
@classmethod
def from_config(cls, config: dict[str, Any]) -> "QuarkConfig":
export_config = config.get("export")
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
index 007e78e68d5cd..33e9f9806b27e 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
@@ -10,6 +10,7 @@ import torch
import torch.nn.functional as F
from vllm import envs
+from vllm._aiter_ops import rocm_aiter_ops
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
dequant_mxfp4,
@@ -49,7 +50,10 @@ def is_rocm_aiter_fp4_asm_gemm_enabled() -> bool:
try:
from aiter.ops.shuffle import shuffle_weight
- from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
+ from aiter.ops.triton.gemm_afp4wfp4 import (
+ gemm_afp4wfp4,
+ gemm_afp4wfp4_preshuffled_weight_scales,
+ )
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from vllm.utils.torch_utils import direct_register_custom_op
@@ -66,23 +70,56 @@ try:
x_scales: torch.Tensor | None = None,
) -> torch.Tensor:
M = x.shape[0]
+ N = weight.shape[0]
+ K = weight.shape[1]
if rocm_use_aiter_fp4_asm_gemm:
- if x_scales is None:
- # use hip quant kernel for performance
- x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
+ if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K):
+ if x_scales is None:
+ # use hip quant kernel for performance
+ if M >= 32:
+ x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
+ else:
+ x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False)
+ else:
+ x_q = x
+ x_s = x_scales
+
+ if M >= 32:
+ x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1)
+ else:
+ x_s = x_s[:M, ...].view(torch.uint8)
+
+ y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
+ gemm_afp4wfp4_preshuffled_weight_scales(
+ x_q.view(torch.uint8),
+ weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
+ x_s,
+ weight_scale.view(torch.uint8).view(
+ weight_scale.shape[0] // 32, -1
+ ),
+ out_dtype,
+ y,
+ )
else:
- x_q = x
- x_s = x_scales
+ if x_scales is None:
+ # use hip quant kernel for performance
+ x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
+ else:
+ x_q = x
+ x_s = x_scales
- # 32 alignment is enough for dim0 padding of output for
- # gemm_a4w4 kernel
- y = torch.empty(
- (M + 31) // 32 * 32, weight.shape[0], device=x_q.device, dtype=out_dtype
- )
+ # 32 alignment is enough for dim0 padding of output for
+ # gemm_a4w4 kernel
+ y = torch.empty(
+ (M + 31) // 32 * 32,
+ weight.shape[0],
+ device=x_q.device,
+ dtype=out_dtype,
+ )
- gemm_a4w4(
- x_q, weight, x_s, weight_scale.view(x_s.dtype), y, bpreshuffle=True
- )
+ gemm_a4w4(
+ x_q, weight, x_s, weight_scale.view(x_s.dtype), y, bpreshuffle=True
+ )
return y[:M]
else:
if x_scales is None:
diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=10240,K=5120,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=10240,K=5120,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
new file mode 100644
index 0000000000000..6b2c1dc1312bf
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/configs/N=10240,K=5120,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
@@ -0,0 +1,146 @@
+{
+ "1": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 5
+ },
+ "2": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "4": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "8": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "16": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "24": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "32": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "48": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 256,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "64": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "96": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 256,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "128": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "256": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "512": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "1024": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "1536": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "2048": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "3072": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4096": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 4
+ }
+}
diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=5120,K=25600,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=5120,K=25600,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
new file mode 100644
index 0000000000000..b0eaf02a541ad
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/configs/N=5120,K=25600,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
@@ -0,0 +1,146 @@
+{
+ "1": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "2": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "4": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "8": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "16": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "24": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "32": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "48": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "64": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "96": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "128": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "256": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "512": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "1024": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "1536": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "2048": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "3072": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "4096": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 256,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ }
+}
diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=5120,K=8192,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=5120,K=8192,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
new file mode 100644
index 0000000000000..4cd357d5086ca
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/configs/N=5120,K=8192,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
@@ -0,0 +1,146 @@
+{
+ "1": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "2": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "8": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "16": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "24": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "32": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "48": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "64": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "96": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 4
+ },
+ "128": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "256": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 2
+ },
+ "512": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "1024": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "1536": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "2048": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "3072": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4096": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 8,
+ "num_stages": 3
+ }
+}
diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=51200,K=5120,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=51200,K=5120,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
new file mode 100644
index 0000000000000..ca2179ddf3d2f
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/configs/N=51200,K=5120,device_name=NVIDIA_L40S,dtype=fp8_w8a8,block_shape=[128,128].json
@@ -0,0 +1,146 @@
+{
+ "1": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "2": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 5
+ },
+ "8": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 1,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "16": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 5
+ },
+ "24": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "32": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "48": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "64": {
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "96": {
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 4,
+ "num_stages": 3
+ },
+ "128": {
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 32,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 32,
+ "num_warps": 4,
+ "num_stages": 5
+ },
+ "256": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "512": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "1024": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "1536": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "2048": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 4
+ },
+ "3072": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 16,
+ "num_warps": 8,
+ "num_stages": 3
+ },
+ "4096": {
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "BLOCK_SIZE_K": 128,
+ "GROUP_SIZE_M": 64,
+ "num_warps": 8,
+ "num_stages": 3
+ }
+}
diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
index fdf330329e20c..eda40657b1e39 100644
--- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
@@ -9,6 +9,10 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEQuantConfig,
+ RoutingMethodType,
+)
+from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (
+ FlashInferCuteDSLExperts,
)
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
FlashInferExperts,
@@ -17,10 +21,14 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize im
create_flashinfer_prepare_finalize,
)
from vllm.platforms import current_platform
-from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+from vllm.utils.flashinfer import (
+ has_flashinfer_cutedsl_grouped_gemm_nt_masked,
+ has_flashinfer_cutlass_fused_moe,
+)
__all__ = [
"is_flashinfer_fp4_cutlass_moe_available",
+ "is_flashinfer_fp4_cutedsl_moe_available",
"reorder_w1w3_to_w3w1",
"build_flashinfer_fp4_cutlass_moe_prepare_finalize",
]
@@ -36,6 +44,16 @@ def is_flashinfer_fp4_cutlass_moe_available() -> bool:
)
+def is_flashinfer_fp4_cutedsl_moe_available() -> bool:
+ """Return ``True`` when FlashInfer CUTEDSL NV-FP4 kernels can be used."""
+ return (
+ envs.VLLM_USE_FLASHINFER_MOE_FP4
+ and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
+ and current_platform.is_cuda()
+ and current_platform.is_device_capability(100)
+ )
+
+
def reorder_w1w3_to_w3w1(
weight: torch.Tensor, scale: torch.Tensor, dim: int = -2
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -72,18 +90,244 @@ def select_nvfp4_gemm_impl(
"""Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""
if allow_flashinfer:
- return FlashInferExperts(
- out_dtype=moe.in_dtype,
- quant_config=moe_quant_config,
- ep_rank=moe.moe_parallel_config.ep_rank,
- ep_size=moe.moe_parallel_config.ep_size,
- tp_rank=moe.moe_parallel_config.tp_rank,
- tp_size=moe.moe_parallel_config.tp_size,
- use_dp=moe.moe_parallel_config.dp_size > 1,
- )
+ if envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm":
+ return FlashInferCuteDSLExperts(
+ out_dtype=moe.in_dtype,
+ quant_config=moe_quant_config,
+ )
+ elif envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput":
+ return FlashInferExperts(
+ out_dtype=moe.in_dtype,
+ quant_config=moe_quant_config,
+ ep_rank=moe.moe_parallel_config.ep_rank,
+ ep_size=moe.moe_parallel_config.ep_size,
+ tp_rank=moe.moe_parallel_config.tp_rank,
+ tp_size=moe.moe_parallel_config.tp_size,
+ use_dp=moe.moe_parallel_config.dp_size > 1,
+ )
# native cutlass experts currently don't support DP; TP case won't call this
raise ValueError(
"CutlassExpertsFp4 doesn't support DP. Use flashinfer CUTLASS "
"Fused MoE backend instead (set VLLM_USE_FLASHINFER_MOE_FP4=1)"
)
+
+
+def prepare_static_weights_for_trtllm_fp4_moe(
+ # args_dequant,
+ # args,
+ gemm1_weights,
+ gemm2_weights,
+ gemm1_scales_linear_fp4_bytes,
+ gemm2_scales_linear_fp4_bytes,
+ hidden_size,
+ intermediate_size,
+ num_experts,
+):
+ from flashinfer import nvfp4_block_scale_interleave
+ from flashinfer.fused_moe.core import (
+ _maybe_get_cached_w3_w1_permute_indices,
+ get_w2_permute_indices_with_cache,
+ )
+
+ _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
+ """Prepare quantized weights for kernel (done offline with weights)."""
+ epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
+
+ # Convert quantized weights to proper formats
+ gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
+ num_experts, 2 * intermediate_size, hidden_size // 2
+ ) # packed fp4
+ gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
+ torch.float8_e4m3fn
+ ).reshape(
+ num_experts, 2 * intermediate_size, hidden_size // 16
+ ) # fp8 scaling factors
+
+ gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
+ num_experts, hidden_size, intermediate_size // 2
+ ) # packed fp4
+ gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(
+ torch.float8_e4m3fn
+ ).reshape(num_experts, hidden_size, intermediate_size // 16) # fp8 scaling factors
+
+ gemm1_weights_fp4_shuffled = []
+ gemm1_scales_fp4_shuffled = []
+ gemm2_weights_fp4_shuffled = []
+ gemm2_scales_fp4_shuffled = []
+ for i in range(num_experts):
+ # Calculate the permute indices for the following:
+ # 1. Reorder rows of W1 and scales for fused gated activation
+ # 2. Shuffle weights and scaling factors for transposed mma output
+ # for both w3_w1 and w2 weights and scale factors
+ permute_indices = _maybe_get_cached_w3_w1_permute_indices(
+ _cache_permute_indices,
+ gemm1_weights_fp4[i].view(torch.uint8),
+ epilogue_tile_m,
+ )
+ gemm1_weights_fp4_shuffled.append(
+ gemm1_weights_fp4[i]
+ .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]
+ .contiguous()
+ )
+
+ permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
+ _cache_permute_indices,
+ gemm1_scales_linear_fp4[i].view(torch.uint8),
+ epilogue_tile_m,
+ num_elts_per_sf=16,
+ )
+ gemm1_scales_fp4_shuffled.append(
+ nvfp4_block_scale_interleave(
+ gemm1_scales_linear_fp4[i]
+ .view(torch.uint8)[
+ permute_sf_indices.to(gemm1_scales_linear_fp4.device)
+ ]
+ .contiguous()
+ )
+ )
+
+ permute_indices = get_w2_permute_indices_with_cache(
+ _cache_permute_indices,
+ gemm2_weights_fp4[i].view(torch.uint8),
+ epilogue_tile_m,
+ )
+ gemm2_weights_fp4_shuffled.append(
+ gemm2_weights_fp4[i]
+ .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]
+ .contiguous()
+ )
+
+ permute_sf_indices = get_w2_permute_indices_with_cache(
+ _cache_permute_indices,
+ gemm2_scales_linear_fp4[i].view(torch.uint8),
+ epilogue_tile_m,
+ num_elts_per_sf=16,
+ )
+ gemm2_scales_fp4_shuffled.append(
+ nvfp4_block_scale_interleave(
+ gemm2_scales_linear_fp4[i]
+ .view(torch.uint8)[
+ permute_sf_indices.to(gemm2_scales_linear_fp4.device)
+ ]
+ .contiguous()
+ )
+ )
+
+ # Stack weights for all experts
+ gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
+ gemm1_scales_fp4_shuffled = (
+ torch.stack(gemm1_scales_fp4_shuffled)
+ .view(torch.float8_e4m3fn)
+ .reshape(num_experts, 2 * intermediate_size, hidden_size // 16)
+ )
+
+ gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)
+ gemm2_scales_fp4_shuffled = (
+ torch.stack(gemm2_scales_fp4_shuffled)
+ .view(torch.float8_e4m3fn)
+ .reshape(num_experts, hidden_size, intermediate_size // 16)
+ )
+ return (
+ gemm1_weights_fp4_shuffled,
+ gemm1_scales_fp4_shuffled,
+ gemm2_weights_fp4_shuffled,
+ gemm2_scales_fp4_shuffled,
+ )
+
+
+def flashinfer_trtllm_fp4_moe(
+ layer: torch.nn.Module,
+ x: torch.Tensor,
+ router_logits: torch.Tensor,
+ top_k: int,
+ global_num_experts: int,
+ num_expert_group: int | None,
+ topk_group: int | None,
+ custom_routing_function: object | None,
+ e_score_correction_bias: torch.Tensor | None,
+) -> torch.Tensor:
+ """
+ Apply FlashInfer TensorRT-LLM FP4 MoE kernel.
+
+ Args:
+ layer: The MoE layer with weights and scales
+ x: Input tensor
+ router_logits: Router logits for expert selection
+ top_k: Number of experts to select per token
+ global_num_experts: Total number of experts across all ranks
+ num_expert_group: Number of expert groups (for grouped routing)
+ topk_group: Top-k within each group
+ custom_routing_function: Custom routing function (e.g., Llama4)
+ e_score_correction_bias: Optional routing bias correction
+
+ Returns:
+ Output tensor from the MoE layer
+ """
+ import flashinfer
+
+ from vllm.model_executor.models.llama4 import Llama4MoE
+
+ # Quantize input to FP4
+ a1_gscale = layer.w13_input_scale_quant
+ (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(
+ x,
+ a1_gscale,
+ is_sf_swizzled_layout=False,
+ )
+
+ # Determine routing method type
+ use_llama4_routing = custom_routing_function is Llama4MoE.custom_routing_function
+ routing_method_type = layer.routing_method_type
+ if use_llama4_routing:
+ routing_method_type = flashinfer.RoutingMethodType.Llama4
+
+ # Prepare routing bias
+ routing_bias = e_score_correction_bias
+ if routing_bias is not None:
+ routing_bias = routing_bias.to(torch.bfloat16)
+
+ router_logits = (
+ router_logits.to(torch.float32)
+ if routing_method_type == RoutingMethodType.DeepSeekV3
+ else router_logits
+ )
+
+ # Call TRT-LLM FP4 block-scale MoE kernel
+ out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
+ routing_logits=router_logits,
+ routing_bias=routing_bias,
+ hidden_states=hidden_states_fp4,
+ hidden_states_scale=hidden_states_scale_linear_fp4.view(
+ torch.float8_e4m3fn
+ ).flatten(),
+ gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,
+ gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(
+ torch.float8_e4m3fn
+ ),
+ gemm1_bias=None,
+ gemm1_alpha=None,
+ gemm1_beta=None,
+ gemm1_clamp_limit=None,
+ gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,
+ gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(
+ torch.float8_e4m3fn
+ ),
+ gemm2_bias=None,
+ output1_scale_scalar=layer.g1_scale_c.data,
+ output1_scale_gate_scalar=layer.g1_alphas.data,
+ output2_scale_scalar=layer.g2_alphas.data,
+ num_experts=global_num_experts,
+ top_k=top_k,
+ n_group=num_expert_group if num_expert_group is not None else 0,
+ topk_group=topk_group if topk_group is not None else 0,
+ intermediate_size=layer.intermediate_size_per_partition,
+ local_expert_offset=layer.ep_rank * layer.local_num_experts,
+ local_num_experts=layer.local_num_experts,
+ routed_scaling_factor=None,
+ tile_tokens_dim=None,
+ routing_method_type=routing_method_type,
+ do_finalize=True,
+ )[0]
+
+ return out
diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
index f22e17945d1f6..eef7a0896c375 100644
--- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
@@ -25,6 +25,7 @@ logger = init_logger(__name__)
class FlashinferMoeBackend(Enum):
TENSORRT_LLM = "TensorRT-LLM"
CUTLASS = "CUTLASS"
+ CUTEDSL = "CUTEDSL"
def calculate_tile_tokens_dim(num_tokens, top_k, num_experts):
@@ -273,19 +274,31 @@ def flashinfer_cutlass_moe_fp8(
def get_flashinfer_moe_backend() -> FlashinferMoeBackend:
- flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND
- # Prefer CUTLASS on SM90 to cover both SM90/SM100 generations
- if flashinfer_moe_backend == "throughput" or current_platform.is_device_capability(
- 90
- ):
- return FlashinferMoeBackend.CUTLASS
- elif flashinfer_moe_backend == "latency":
- return FlashinferMoeBackend.TENSORRT_LLM
+ backend_map = {
+ "throughput": FlashinferMoeBackend.CUTLASS,
+ "latency": FlashinferMoeBackend.TENSORRT_LLM,
+ "masked_gemm": FlashinferMoeBackend.CUTEDSL,
+ }
+
+ flashinfer_moe_backend = envs.VLLM_FLASHINFER_MOE_BACKEND
+ if flashinfer_moe_backend in backend_map:
+ if (
+ flashinfer_moe_backend == "latency"
+ and not current_platform.is_device_capability(100)
+ ):
+ logger.info_once(
+ "Flashinfer TRTLLM MOE backend is only supported on "
+ "SM100 and later, using CUTLASS backend instead",
+ scope="local",
+ )
+ return FlashinferMoeBackend.CUTLASS
+ return backend_map[flashinfer_moe_backend]
+ elif current_platform.is_device_capability(90):
+ return FlashinferMoeBackend.CUTLASS
- allowed_backends = ["throughput", "latency"]
raise ValueError(
- f"Unknown flashinfer moe backend: {flashinfer_moe_backend}"
- f" expected one of {allowed_backends}"
+ f"Unknown flashinfer moe backend: {flashinfer_moe_backend!r}. "
+ f"Expected one of {list(backend_map.keys())}."
)
diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
index 34a31bcf6a747..d0c8b3d1a3093 100644
--- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
@@ -8,6 +8,7 @@ import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import triton
+from vllm.utils.import_utils import has_triton_kernels
from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
logger = init_logger(__name__)
@@ -15,6 +16,7 @@ logger = init_logger(__name__)
def _swizzle_mxfp4(quant_tensor, scale, num_warps):
"""weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
+ assert has_triton_kernels()
import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
from triton_kernels.numerics import InFlexData
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
@@ -37,15 +39,15 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
value_layout = StridedLayout
scale_layout = StridedLayout
elif current_platform.is_rocm():
- from triton_kernels.tensor_details.layout import (
- GFX950MXScaleLayout,
- StridedLayout,
- )
-
from vllm.platforms.rocm import on_gfx950
value_layout = StridedLayout
- scale_layout = GFX950MXScaleLayout if on_gfx950() else StridedLayout
+ if on_gfx950():
+ from triton_kernels.tensor_details.layout import GFX950MXScaleLayout
+
+ scale_layout = GFX950MXScaleLayout
+ else:
+ scale_layout = StridedLayout
else:
value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
mx_axis=1
diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py b/vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
index c3f26cc774118..44c5b027daf4f 100644
--- a/vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
+++ b/vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
@@ -5,6 +5,7 @@ from dataclasses import dataclass
import vllm.envs as envs
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+ is_flashinfer_fp4_cutedsl_moe_available,
is_flashinfer_fp4_cutlass_moe_available,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
@@ -32,7 +33,10 @@ def detect_nvfp4_moe_support(class_name: str = "") -> NvFp4Support:
"""Detect platform support for NV-FP4 fused-MoE path"""
cutlass_supported = cutlass_fp4_supported()
- allow_flashinfer = cutlass_supported and is_flashinfer_fp4_cutlass_moe_available()
+ allow_flashinfer = cutlass_supported and (
+ is_flashinfer_fp4_cutlass_moe_available()
+ or is_flashinfer_fp4_cutedsl_moe_available()
+ )
if allow_flashinfer:
_logger.info_once(
diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py
index 56c165f9c041a..152d9401b8e94 100644
--- a/vllm/model_executor/layers/rotary_embedding/__init__.py
+++ b/vllm/model_executor/layers/rotary_embedding/__init__.py
@@ -26,23 +26,23 @@ def get_rope(
head_size: int,
rotary_dim: int,
max_position: int,
- base: float,
is_neox_style: bool = True,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
dtype: torch.dtype | None = None,
partial_rotary_factor: float = 1.0,
dual_chunk_attention_config: dict[str, Any] | None = None,
) -> RotaryEmbedding:
if dtype is None:
dtype = torch.get_default_dtype()
- if rope_scaling is not None:
+ if rope_parameters is not None:
# Transforms every value that is a list into a tuple for caching calls
- rope_scaling_tuple = {
- k: tuple(v) if isinstance(v, list) else v for k, v in rope_scaling.items()
+ rope_parameters_tuple = {
+ k: tuple(v) if isinstance(v, list) else v
+ for k, v in rope_parameters.items()
}
- rope_scaling_args = tuple(rope_scaling_tuple.items())
+ rope_parameters_args = tuple(rope_parameters_tuple.items())
else:
- rope_scaling_args = None
+ rope_parameters_args = None
if dual_chunk_attention_config is not None:
dual_chunk_attention_tuple = {
@@ -60,15 +60,15 @@ def get_rope(
head_size,
rotary_dim,
max_position,
- base,
is_neox_style,
- rope_scaling_args,
+ rope_parameters_args,
dual_chunk_attention_args,
dtype,
)
if key in _ROPE_DICT:
return _ROPE_DICT[key]
+ base = rope_parameters["rope_theta"] if rope_parameters else 10000
if dual_chunk_attention_config is not None:
extra_kwargs = {
k: v
@@ -84,18 +84,18 @@ def get_rope(
dtype,
**extra_kwargs,
)
- elif not rope_scaling:
+ elif not rope_parameters:
rotary_emb = RotaryEmbedding(
head_size, rotary_dim, max_position, base, is_neox_style, dtype
)
else:
- scaling_type = rope_scaling["rope_type"]
+ scaling_type = rope_parameters["rope_type"]
if scaling_type == "llama3":
- scaling_factor = rope_scaling["factor"]
- low_freq_factor = rope_scaling["low_freq_factor"]
- high_freq_factor = rope_scaling["high_freq_factor"]
- original_max_position = rope_scaling["original_max_position_embeddings"]
+ scaling_factor = rope_parameters["factor"]
+ low_freq_factor = rope_parameters["low_freq_factor"]
+ high_freq_factor = rope_parameters["high_freq_factor"]
+ original_max_position = rope_parameters["original_max_position_embeddings"]
rotary_emb = Llama3RotaryEmbedding(
head_size,
rotary_dim,
@@ -113,7 +113,7 @@ def get_rope(
head_size, rotary_dim, max_position, base, is_neox_style, dtype
)
elif scaling_type == "default":
- if "mrope_section" in rope_scaling:
+ if "mrope_section" in rope_parameters:
rotary_emb = MRotaryEmbedding(
head_size,
rotary_dim,
@@ -121,8 +121,8 @@ def get_rope(
base,
is_neox_style,
dtype,
- mrope_section=rope_scaling["mrope_section"],
- mrope_interleaved=rope_scaling.get("mrope_interleaved", False),
+ mrope_section=rope_parameters["mrope_section"],
+ mrope_interleaved=rope_parameters.get("mrope_interleaved", False),
)
else:
rotary_emb = RotaryEmbedding(
@@ -134,7 +134,7 @@ def get_rope(
dtype,
)
elif scaling_type == "linear":
- scaling_factor = rope_scaling["factor"]
+ scaling_factor = rope_parameters["factor"]
rotary_emb = LinearScalingRotaryEmbedding(
head_size,
rotary_dim,
@@ -145,8 +145,8 @@ def get_rope(
dtype,
)
elif scaling_type == "ntk":
- scaling_factor = rope_scaling["factor"]
- mixed_b = rope_scaling.get("mixed_b", None)
+ scaling_factor = rope_parameters["factor"]
+ mixed_b = rope_parameters.get("mixed_b")
rotary_emb = NTKScalingRotaryEmbedding(
head_size,
rotary_dim,
@@ -158,8 +158,8 @@ def get_rope(
mixed_b,
)
elif scaling_type == "dynamic":
- if "alpha" in rope_scaling:
- scaling_alpha = rope_scaling["alpha"]
+ if "alpha" in rope_parameters:
+ scaling_alpha = rope_parameters["alpha"]
rotary_emb = DynamicNTKAlphaRotaryEmbedding(
head_size,
rotary_dim,
@@ -169,8 +169,8 @@ def get_rope(
scaling_alpha,
dtype,
)
- elif "factor" in rope_scaling:
- scaling_factor = rope_scaling["factor"]
+ elif "factor" in rope_parameters:
+ scaling_factor = rope_parameters["factor"]
rotary_emb = DynamicNTKScalingRotaryEmbedding(
head_size,
rotary_dim,
@@ -185,11 +185,11 @@ def get_rope(
"Dynamic rope scaling must contain either 'alpha' or 'factor' field"
)
elif scaling_type == "yarn":
- scaling_factor = rope_scaling["factor"]
- original_max_position = rope_scaling["original_max_position_embeddings"]
+ scaling_factor = rope_parameters["factor"]
+ original_max_position = rope_parameters["original_max_position_embeddings"]
extra_kwargs = {
k: v
- for k, v in rope_scaling.items()
+ for k, v in rope_parameters.items()
if k
in (
"extrapolation_factor",
@@ -197,9 +197,10 @@ def get_rope(
"beta_fast",
"beta_slow",
"apply_yarn_scaling",
+ "truncate",
)
}
- if "mrope_section" in rope_scaling:
+ if "mrope_section" in rope_parameters:
extra_kwargs.pop("apply_yarn_scaling", None)
rotary_emb = MRotaryEmbedding(
head_size,
@@ -208,8 +209,8 @@ def get_rope(
base,
is_neox_style,
dtype,
- mrope_section=rope_scaling["mrope_section"],
- mrope_interleaved=rope_scaling.get("mrope_interleaved", False),
+ mrope_section=rope_parameters["mrope_section"],
+ mrope_interleaved=rope_parameters.get("mrope_interleaved", False),
scaling_factor=scaling_factor,
**extra_kwargs,
)
@@ -225,12 +226,12 @@ def get_rope(
**extra_kwargs,
)
elif scaling_type == "deepseek_yarn":
- scaling_factor = rope_scaling["factor"]
- original_max_position = rope_scaling["original_max_position_embeddings"]
+ scaling_factor = rope_parameters["factor"]
+ original_max_position = rope_parameters["original_max_position_embeddings"]
# assert max_position == original_max_position * scaling_factor
extra_kwargs = {
k: v
- for k, v in rope_scaling.items()
+ for k, v in rope_parameters.items()
if k
in (
"extrapolation_factor",
@@ -252,12 +253,12 @@ def get_rope(
**extra_kwargs,
)
elif scaling_type == "longrope":
- short_factor = rope_scaling["short_factor"]
- long_factor = rope_scaling["long_factor"]
- original_max_position = rope_scaling["original_max_position_embeddings"]
+ short_factor = rope_parameters["short_factor"]
+ long_factor = rope_parameters["long_factor"]
+ original_max_position = rope_parameters["original_max_position_embeddings"]
extra_kwargs = {
k: v
- for k, v in rope_scaling.items()
+ for k, v in rope_parameters.items()
if k in ("short_mscale", "long_mscale")
}
rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(
diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py
index ce4f40680b0a3..4114b21168cc8 100644
--- a/vllm/model_executor/layers/rotary_embedding/base.py
+++ b/vllm/model_executor/layers/rotary_embedding/base.py
@@ -83,6 +83,11 @@ class RotaryEmbeddingBase(CustomOp):
):
self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype)
+ def get_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]:
+ cos_sin = self.cos_sin_cache[:seqlen]
+ cos, sin = cos_sin.chunk(2, dim=-1)
+ return cos, sin
+
class RotaryEmbedding(RotaryEmbeddingBase):
def __init__(
diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py
index 196533b617959..13f8d15cc0f72 100644
--- a/vllm/model_executor/layers/rotary_embedding/common.py
+++ b/vllm/model_executor/layers/rotary_embedding/common.py
@@ -117,13 +117,13 @@ def yarn_find_correction_range(
dim: int,
base: float = 10000,
max_position_embeddings: int = 2048,
-) -> tuple[int, int]:
- low = math.floor(
- yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
- )
- high = math.ceil(
- yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
- )
+ truncate: bool = True,
+) -> tuple[float | int, float | int]:
+ low = yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
+ high = yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
+ if truncate:
+ low = math.floor(low)
+ high = math.ceil(high)
return max(low, 0), min(high, dim - 1) # Clamp values just in case
diff --git a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
index ff46ad74b302e..f01ca1e231211 100644
--- a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
+++ b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
@@ -28,12 +28,14 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
beta_fast: int = 32,
beta_slow: int = 1,
apply_yarn_scaling: bool = True,
+ truncate: bool = True,
) -> None:
self.scaling_factor = scaling_factor
self.extrapolation_factor = extrapolation_factor
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
+ self.truncate = truncate
# Get n-d magnitude scaling corrected for interpolation
self.mscale = (
float(yarn_get_mscale(self.scaling_factor) * attn_factor)
@@ -57,6 +59,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
self.rotary_dim,
self.base,
self.max_position_embeddings,
+ self.truncate,
)
# Get n-d rotational scaling corrected for extrapolation
inv_freq_mask = (
diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py
index 301f2d00bf404..052d2cfc1099e 100644
--- a/vllm/model_executor/model_loader/__init__.py
+++ b/vllm/model_executor/model_loader/__init__.py
@@ -30,6 +30,7 @@ logger = init_logger(__name__)
# if a new load format is added here
LoadFormats = Literal[
"auto",
+ "hf",
"bitsandbytes",
"dummy",
"fastsafetensors",
@@ -45,6 +46,7 @@ LoadFormats = Literal[
]
_LOAD_FORMAT_TO_MODEL_LOADER: dict[str, type[BaseModelLoader]] = {
"auto": DefaultModelLoader,
+ "hf": DefaultModelLoader,
"bitsandbytes": BitsAndBytesModelLoader,
"dummy": DummyModelLoader,
"fastsafetensors": DefaultModelLoader,
diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py
index c06ac550a94ae..7401a7a0e2dbb 100644
--- a/vllm/model_executor/model_loader/default_loader.py
+++ b/vllm/model_executor/model_loader/default_loader.py
@@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import (
fastsafetensors_weights_iterator,
filter_duplicate_safetensors_files,
filter_files_not_needed_for_inference,
+ get_quant_config,
maybe_download_from_modelscope,
multi_thread_pt_weights_iterator,
multi_thread_safetensors_weights_iterator,
@@ -30,6 +31,7 @@ from vllm.model_executor.model_loader.weight_utils import (
safetensors_weights_iterator,
)
from vllm.platforms import current_platform
+from vllm.transformers_utils.config import list_filtered_repo_files
logger = init_logger(__name__)
@@ -95,8 +97,25 @@ class DefaultModelLoader(BaseModelLoader):
load_format = self.load_config.load_format
use_safetensors = False
index_file = SAFE_WEIGHTS_INDEX_NAME
- # Some quantized models use .pt files for storing the weights.
+
+ # First check for 'auto' format that mistral files format are present.
+ # This is to load mistral models with official format by default.
if load_format == "auto":
+ load_format = (
+ "mistral"
+ if len(
+ list_filtered_repo_files(
+ model_name_or_path=model_name_or_path,
+ allow_patterns=["consolidated*.safetensors"],
+ revision=revision,
+ )
+ )
+ > 0
+ else "hf"
+ )
+
+ # Some quantized models use .pt files for storing the weights.
+ if load_format == "hf":
allow_patterns = ["*.safetensors", "*.bin"]
elif load_format == "safetensors" or load_format == "fastsafetensors":
use_safetensors = True
@@ -273,42 +292,17 @@ class DefaultModelLoader(BaseModelLoader):
)
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
- if model_config.quantization == "torchao" and torchao_version_at_least(
- "0.14.0"
- ):
- self.load_config.safetensors_load_strategy = "torchao"
+ if model_config.quantization == "torchao":
+ quant_config = get_quant_config(model_config, self.load_config)
+ if (
+ hasattr(quant_config, "is_checkpoint_torchao_serialized")
+ and quant_config.is_checkpoint_torchao_serialized
+ and torchao_version_at_least("0.15.0")
+ ):
+ self.load_config.safetensors_load_strategy = "torchao"
+
weights_to_load = {name for name, _ in model.named_parameters()}
-
- # if we don't have `model.weight_metadata_and_attr_saved` defined and
- # set to True, it means that this is either offline quantization case
- # or the first run of online quantization
- # see online_quantization.py for detailed notes
- offline_quantization_or_first_run_of_online_quantization = not getattr(
- model, "weight_metadata_and_attr_saved", False
- )
-
- if model_config.quantization is None:
- # model is not quantized
- loaded_weights = model.load_weights(
- self.get_all_weights(model_config, model)
- )
- elif offline_quantization_or_first_run_of_online_quantization:
- # case 1: offline quantized checkpoint
- # case 2: Step I1 first run of weight loading with
- # online quantization
- # see online_quantization.py for detailed notes
- loaded_weights = model.load_weights(
- self.get_all_weights(model_config, model)
- )
- else:
- # to avoid circular dependency
- from vllm.model_executor.model_loader.online_quantization import (
- load_weights_and_online_quantize,
- )
-
- # subsequent runs of weight loading with online
- # quantization
- loaded_weights = load_weights_and_online_quantize(self, model, model_config)
+ loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
self.counter_after_loading_weights = time.perf_counter()
logger.info_once(
diff --git a/vllm/model_executor/model_loader/gguf_loader.py b/vllm/model_executor/model_loader/gguf_loader.py
index 7db1fc167c4fa..2416836be03c4 100644
--- a/vllm/model_executor/model_loader/gguf_loader.py
+++ b/vllm/model_executor/model_loader/gguf_loader.py
@@ -7,10 +7,11 @@ import gguf
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoModelForImageTextToText
from vllm.config import ModelConfig, VllmConfig
from vllm.config.load import LoadConfig
+from vllm.logger import init_logger
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
from vllm.model_executor.model_loader.utils import (
initialize_model,
@@ -21,8 +22,11 @@ from vllm.model_executor.model_loader.weight_utils import (
get_gguf_weight_type_map,
gguf_quant_weights_iterator,
)
+from vllm.transformers_utils.gguf_utils import detect_gguf_multimodal
from vllm.utils.torch_utils import set_default_torch_dtype
+logger = init_logger(__name__)
+
class GGUFModelLoader(BaseModelLoader):
"""
@@ -67,7 +71,15 @@ class GGUFModelLoader(BaseModelLoader):
https://github.com/ggerganov/ggml/blob/master/docs/gguf.md for details.
"""
config = model_config.hf_config
+ # Get text config to handle both nested (multimodal) and flat
+ # (text-only) config structures. For multimodal models like
+ # Gemma3Config, this returns config.text_config. For text-only
+ # models, this returns config itself.
+ text_config = config.get_text_config()
model_type = config.model_type
+ is_multimodal = (
+ hasattr(config, "vision_config") and config.vision_config is not None
+ )
gguf_to_hf_name_map = {}
# hack: ggufs have a different name than transformers
if model_type == "cohere":
@@ -115,24 +127,167 @@ class GGUFModelLoader(BaseModelLoader):
break
if arch is None:
raise RuntimeError(f"Unknown gguf model_type: {model_type}")
- num_layers = config.num_hidden_layers
- name_map = gguf.get_tensor_name_map(arch, num_layers)
+ text_num_layers = text_config.num_hidden_layers
+ text_name_map = gguf.get_tensor_name_map(arch, text_num_layers)
+
+ if is_multimodal:
+ mm_proj_arch = gguf.MODEL_ARCH.MMPROJ
+ vision_num_layers = config.vision_config.num_hidden_layers
+ vision_name_map = gguf.get_tensor_name_map(mm_proj_arch, vision_num_layers)
+ else:
+ vision_name_map = None
+
+ # Create dummy model to extract parameter names
+ # For multimodal: use AutoModelForImageTextToText to get
+ # language + vision + projector params
+ # For text-only: use AutoModelForCausalLM to get language model params
+ auto_cls = (
+ AutoModelForImageTextToText if is_multimodal else AutoModelForCausalLM
+ )
with torch.device("meta"):
- dummy_model = AutoModelForCausalLM.from_config(
+ dummy_model = auto_cls.from_config(
config, trust_remote_code=model_config.trust_remote_code
)
- state_dict = dummy_model.state_dict()
+ state_dict = dummy_model.state_dict()
+ if hf_checkpoint_map := getattr(
+ dummy_model, "_checkpoint_conversion_mapping", None
+ ):
+
+ def revert_hf_rename(name: str) -> str:
+ for original_name, hf_name in hf_checkpoint_map.items():
+ if hf_name in name:
+ name = name.replace(hf_name, original_name).lstrip("^")
+ return name
+
+ state_dict = {
+ revert_hf_rename(name): tensor for name, tensor in state_dict.items()
+ }
+
+ def find_hf_name_in_tensor_map(hf_name: str) -> str | None:
+ """
+ Map HuggingFace parameter name to GGUF tensor name.
+
+ This function handles the mismatch between HF parameter naming
+ conventions and gguf-py's expected format:
+ 1. Strips 'model.' prefix (common in multimodal models)
+ 2. Converts '_weight' suffix to '.weight' (Gemma3 compatibility)
+ 3. Searches vision_name_map for multimodal parameters
+ 4. Falls back to text_name_map for language model parameters
+
+ Args:
+ hf_name: Full HuggingFace parameter name (e.g.,
+ 'model.multi_modal_projector.mm_soft_emb_norm.weight')
+
+ Returns:
+ GGUF tensor name with suffix (e.g., 'mm.soft_emb_norm.weight')
+ or None if no mapping found
+ """
+ # Strip 'language_model.' prefix for multimodal models - gguf-py
+ # tensor mappings expect parameter names without this prefix.
+ # Note: 'model.' prefix should be KEPT for text-only models as
+ # gguf-py expects it.
+ if hf_name.startswith("language_model."):
+ hf_name = hf_name[15:] # Remove 'language_model.'
+
+ # Parse parameter name and suffix
+ if hf_name.endswith((".weight", ".bias")):
+ base_name, suffix = hf_name.rsplit(".", 1)
+ else:
+ base_name, suffix = hf_name, ""
+ # Handle '_weight' suffix (Gemma3 naming: parameter ends with
+ # '_weight' instead of '.weight')
+ if base_name.endswith("_weight"):
+ base_name = base_name[:-7] # Remove '_weight'
+ suffix = "weight"
+
+ gguf_name = None
+ # Priority 1: Search vision/projector parameters for multimodal models
+ if vision_name_map is not None:
+ gguf_name = vision_name_map.get_name(base_name)
+
+ # Priority 2: Search text backbone parameters
+ if gguf_name is None:
+ gguf_name = text_name_map.get_name(base_name)
+
+ if gguf_name is None:
+ return None
+
+ return gguf_name + "." + suffix
+
+ # Build mapping and track unmapped parameters
+ unmapped_params = []
for hf_name in state_dict:
- name, suffix = hf_name.rsplit(".", 1)
- gguf_name = name_map.get_name(name)
- gguf_to_hf_name_map[f"{gguf_name}.{suffix}"] = hf_name
+ gguf_name_with_suffix = find_hf_name_in_tensor_map(hf_name)
+
+ # Track mapping success
+ if gguf_name_with_suffix is not None:
+ gguf_to_hf_name_map[gguf_name_with_suffix] = hf_name
+ logger.debug("Mapped GGUF %s → HF %s", gguf_name_with_suffix, hf_name)
+ elif hf_name not in gguf_to_hf_name_map.values():
+ # Parameter not in manual overrides either
+ unmapped_params.append(hf_name)
+
+ # All parameters must be mapped: both vision/projector and backbone
+ if unmapped_params:
+ raise RuntimeError(
+ f"Failed to map GGUF parameters "
+ f"({len(unmapped_params)}): "
+ f"{unmapped_params}"
+ )
return gguf_to_hf_name_map
+ def _get_gguf_weight_type(
+ self,
+ model_config: ModelConfig,
+ model_name_or_path: str,
+ gguf_to_hf_name_map: dict[str, str],
+ ) -> dict[str, str]:
+ weight_type_map = get_gguf_weight_type_map(
+ model_config.model, gguf_to_hf_name_map
+ )
+ is_multimodal = hasattr(model_config.hf_config, "vision_config")
+ if is_multimodal:
+ mmproj_file = detect_gguf_multimodal(model_name_or_path)
+ assert mmproj_file is not None, (
+ "Could not find mm_proj file for multimodal GGUF model"
+ )
+ logger.info("Loading extra mm_proj weights from %s...", mmproj_file)
+ mm_proj_weight_type_map = get_gguf_weight_type_map(
+ mmproj_file, gguf_to_hf_name_map
+ )
+ weight_type_map.update(mm_proj_weight_type_map)
+ return weight_type_map
+
def _get_weights_iterator(
- self, model_name_or_path: str, gguf_to_hf_name_map: dict[str, str]
+ self,
+ model_config: ModelConfig,
+ model_name_or_path: str,
+ gguf_to_hf_name_map: dict[str, str],
) -> Generator[tuple[str, torch.Tensor], None, None]:
- return gguf_quant_weights_iterator(model_name_or_path, gguf_to_hf_name_map)
+ """
+ Iterate over GGUF model weights, loading from both main model file and
+ mmproj.gguf for multimodal Gemma3 models.
+
+ For Gemma3 multimodal GGUF models:
+ - Main file (gemma-3-*.gguf): Language model weights (model.*)
+ - mmproj file (mmproj*.gguf): Vision tower + projector weights (v.*, mm.*)
+
+ Yields:
+ Tuples of (parameter_name, tensor) for all model weights
+ """
+ hf_config = model_config.hf_config
+ is_multimodal = hasattr(hf_config, "vision_config")
+
+ if is_multimodal:
+ # Load mm_proj (mm_encoder + projector) for multimodal weights
+ mmproj_file = detect_gguf_multimodal(model_name_or_path)
+ assert mmproj_file is not None, (
+ "Could not find mm_proj file for multimodal GGUF model"
+ )
+ yield from gguf_quant_weights_iterator(mmproj_file, gguf_to_hf_name_map)
+
+ yield from gguf_quant_weights_iterator(model_name_or_path, gguf_to_hf_name_map)
def download_model(self, model_config: ModelConfig) -> None:
self._prepare_weights(model_config.model)
@@ -141,7 +296,7 @@ class GGUFModelLoader(BaseModelLoader):
local_model_path = self._prepare_weights(model_config.model)
gguf_weights_map = self._get_gguf_weights_map(model_config)
model.load_weights(
- self._get_weights_iterator(local_model_path, gguf_weights_map)
+ self._get_weights_iterator(model_config, local_model_path, gguf_weights_map)
)
def load_model(
@@ -156,14 +311,19 @@ class GGUFModelLoader(BaseModelLoader):
):
model_config.hf_config.update({"tie_word_embeddings": True})
- weight_type_map = get_gguf_weight_type_map(model_config.model, gguf_weights_map)
-
+ weight_type_map = self._get_gguf_weight_type(
+ model_config, local_model_path, gguf_weights_map
+ )
# filter out unquantized modules to skip
unquant_names = [
name.removesuffix(".weight")
for name, weight_type in weight_type_map.items()
- if weight_type == "F32" and name.endswith(".weight")
+ if weight_type in ("F32", "F16", "BF16") and name.endswith(".weight")
]
+ logger.debug(
+ "GGUF unquantized modules: %s",
+ unquant_names,
+ )
vllm_config.quant_config.unquantized_modules.extend(unquant_names)
target_device = torch.device(device_config.device)
diff --git a/vllm/model_executor/model_loader/online_quantization.py b/vllm/model_executor/model_loader/online_quantization.py
index 890dd7231a0e1..f330af85bbe8b 100644
--- a/vllm/model_executor/model_loader/online_quantization.py
+++ b/vllm/model_executor/model_loader/online_quantization.py
@@ -2,13 +2,13 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import types
+from collections.abc import Iterable
import torch
from torch import nn
from vllm.config import ModelConfig
from vllm.logger import init_logger
-from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
from vllm.model_executor.model_loader.utils import process_weights_after_loading
logger = init_logger(__name__)
@@ -56,6 +56,9 @@ logger = init_logger(__name__)
# R4. quantize weights (by calling process_weights_after_loading),
# also set `process_weights_after_loading_already_called` to
# True to stop it from running again
+# R5. (workaround for cudagraph), we restore the weight params to original quantized
+# weights params, and use original_weight_param.copy_(updated_weight_param) so that
+# the weight update work well with cudagraph
# process_weights_after_loading (if called):
# this will be skipped since it's already ran in
# load_weights
@@ -69,14 +72,6 @@ def maybe_save_metadata_and_attributes_for_weight_reloading(
if model_config.quantization != "torchao":
return
- if getattr(model, "process_weights_after_loading_already_called", False):
- # In case `process_weights_after_loading` is called multiple times
- # we'll skip it at later times
- logger.warning(
- "process_weights_after_loading already called for model %s", model
- )
- return
-
from vllm.model_executor.model_loader.weight_utils import get_quant_config
quant_config = get_quant_config(model_config, None)
@@ -137,6 +132,7 @@ def maybe_save_metadata_and_attributes_for_weight_reloading(
else:
model.recorded_weight_attr[name][key] = attr
# mark the metadata and attributes saved so we don't run it again
+ model._model_config = model_config
model.weight_metadata_and_attr_saved = True
@@ -148,77 +144,132 @@ def _bond_method_to_cls(func, obj):
return types.MethodType(func, obj)
-def load_weights_and_online_quantize(
- model_loader: DefaultModelLoader, model: nn.Module, model_config: ModelConfig
-) -> set[str]:
+def support_quantized_model_reload_from_hp_weights(original_load_weights):
+ """Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
+ reloading high precision (bfloat16/float16/float32) weight for an already quantized
+ model, this involves restoring the weights to a high precision weights and
+ then online quantize the weights
+ """
# online quantization, right now only enabled for
# torchao
- # R1, R2, R3, R4 in the Notes
+ # R1, R2, R3, R4, R5 in the Notes
- # TODO: Add fp8 support
- assert model_config.quantization == "torchao", (
- "online quantization is only enabled for torchao currently"
- )
- # TODO: use create_weights to restore the weights to original state
+ def patched_model_load_weights(
+ auto_weight_loader, weights: Iterable[tuple[str, torch.Tensor]], *, mapper=None
+ ) -> set[str]:
+ model = auto_weight_loader.module
+ offline_quantization_or_first_run_of_online_quantization = not getattr(
+ model, "weight_metadata_and_attr_saved", False
+ )
- # Step R1: First restore the quantized weights to original bfloat16
- # weights, with original metadata (shape, dtype, device)
- # and attributes, so that bfloat16 weights can be loaded properly
- existing_param_names = dict(model.named_parameters(remove_duplicate=False)).keys()
- named_modules = dict(model.named_modules(remove_duplicate=False))
- model_device = None
+ # if we don't have `model.weight_metadata_and_attr_saved` defined and
+ # set to True, it means that this is either offline quantization case
+ # or the first run of online quantization
+ # see Notes in this file for more details
+ if offline_quantization_or_first_run_of_online_quantization:
+ # case 1: offline quantized checkpoint
+ # case 2: Step I1 first run of weight loading with
+ # online quantization
+ return original_load_weights(auto_weight_loader, weights, mapper=mapper)
- # Step R2: recover the parameter to the state before first loading
- for name, d in model.original_weights_rebuild_keys.items():
- _shape = d["shape"]
- _dtype = d["dtype"]
- _device = d["device"]
+ model_config = model._model_config
+
+ # TODO: Add fp8 support
+ assert model_config.quantization == "torchao", (
+ "online quantization is only enabled for torchao currently"
+ )
+ # TODO: use create_weights to restore the weights to original state
+
+ # Step R1: First restore the quantized weights to original bfloat16
+ # weights, with original metadata (shape, dtype, device)
+ # and attributes, so that bfloat16 weights can be loaded properly
+ # TODO: maybe set remove_duplicate to True?
+ original_quantized_weight_dict = dict(
+ model.named_parameters(remove_duplicate=False)
+ )
+ named_modules = dict(model.named_modules(remove_duplicate=False))
+ model_device = None
+
+ for name, d in model.original_weights_rebuild_keys.items():
+ _shape = d["shape"]
+ _dtype = d["dtype"]
+ _device = d["device"]
+ if model_device is not None:
+ assert model_device == _device, (
+ "Expecting all weights "
+ "to be in the same device for now, got both: "
+ f"{model_device} and {_device}"
+ )
+ else:
+ model_device = _device
+
+ if name in original_quantized_weight_dict:
+ module_name, weight_name = name.rsplit(".", 1)
+ module = named_modules[module_name]
+ setattr(
+ module,
+ weight_name,
+ torch.nn.Parameter(
+ torch.empty(_shape, dtype=_dtype, device=_device),
+ requires_grad=False,
+ ),
+ )
+
+ # Step R2: recover the weight attributes to the state before first loading
+ # recorded_weight_attr is
+ # {"weight_name": {"weight_attr_key": attr}}
+ # e.g.
+ # {
+ # {
+ # "layer.0.weight": {
+ # "weight_loader": weight_loader_function_object,
+ # "input_dim": 0, ...
+ # },
+ # "layer.1.weight": ...,
+ # }
+ # }
+ for full_weight_name, weight_attr_dict in model.recorded_weight_attr.items():
+ for attr_name, attr in weight_attr_dict.items():
+ module_name, weight_name = full_weight_name.rsplit(".", 1)
+ module = named_modules[module_name]
+ weight = getattr(module, weight_name)
+ if not hasattr(weight, attr_name):
+ setattr(weight, attr_name, _bond_method_to_cls(attr, weight))
+
+ # Step R3: reload bfloat16 / high precision weights
+ updated_params = original_load_weights(
+ auto_weight_loader, weights, mapper=mapper
+ )
+
+ # Step R4: online quantize the weights
+ # manually process weights after loading
+ model.process_weights_after_loading_already_called = False
if model_device is not None:
- assert model_device == _device, (
- "Expecting all weights "
- "to be in the same device for now, got both: "
- f"{model_device} and {_device}"
- )
+ process_weights_after_loading(model, model_config, model_device)
else:
- model_device = _device
-
- if name in existing_param_names:
- module_name, weight_name = name.rsplit(".", 1)
- module = named_modules[module_name]
- setattr(
- module,
- weight_name,
- torch.nn.Parameter(torch.empty(_shape, dtype=_dtype, device=_device)),
+ logger.warning_once(
+ "model_device is None, skip calling process_weights_after_loading"
)
- # recorded_weight_attr is
- # {"weight_name": {"weight_attr_key": attr}}
- # e.g.
- # {
- # {
- # "layer.0.weight": {
- # "weight_loader": weight_loader_function_object,
- # "input_dim": 0, ...
- # },
- # "layer.1.weight": ...,
- # }
- # }
- for full_weight_name, weight_attr_dict in model.recorded_weight_attr.items():
- for attr_name, attr in weight_attr_dict.items():
- module_name, weight_name = full_weight_name.rsplit(".", 1)
- module = named_modules[module_name]
- weight = getattr(module, weight_name)
- if not hasattr(weight, attr_name):
- setattr(weight, attr_name, _bond_method_to_cls(attr, weight))
+ # Step R5 (workaround for cudagraph): restore the original quantized weights
+ # and do a copy_ of the currents weights to the original weights
+ updated_quantized_weights = dict(model.named_parameters(remove_duplicate=False))
+ for name in model.original_weights_rebuild_keys:
+ if name in original_quantized_weight_dict:
+ original_quantized_weight = original_quantized_weight_dict[name]
+ updated_quantized_weight = updated_quantized_weights[name]
- # Step I1: reload bfloat16 / high precision weights
- loaded_weights = model.load_weights(
- model_loader.get_all_weights(model_config, model)
- )
+ module_name, weight_name = name.rsplit(".", 1)
+ module = named_modules[module_name]
+ setattr(module, weight_name, original_quantized_weight)
+ with torch.no_grad():
+ original_quantized_weight.copy_(updated_quantized_weight)
- # Step I2: online quantize the weights
- # manually process weights after loading
- model.process_weights_after_loading_already_called = False
- process_weights_after_loading(model, model_config, model_device)
- model.process_weights_after_loading_already_called = True
- return loaded_weights
+ del original_quantized_weight_dict
+ del named_modules
+ del updated_quantized_weight
+
+ model.process_weights_after_loading_already_called = True
+ return updated_params
+
+ return patched_model_load_weights
diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py
index d94dbd9f06e0b..1538f0c2af655 100644
--- a/vllm/model_executor/model_loader/sharded_state_loader.py
+++ b/vllm/model_executor/model_loader/sharded_state_loader.py
@@ -4,6 +4,7 @@
import collections
import glob
import os
+import time
from collections.abc import Generator
from typing import Any
@@ -132,6 +133,7 @@ class ShardedStateLoader(BaseModelLoader):
f"pre-sharded checkpoints are currently supported!"
)
state_dict = self._filter_subtensors(model.state_dict())
+ counter_before_loading_weights = time.perf_counter()
for key, tensor in self.iterate_over_files(filepaths):
# If loading with LoRA enabled, additional padding may
# be added to certain parameters. We only load into a
@@ -150,6 +152,12 @@ class ShardedStateLoader(BaseModelLoader):
)
param_data.copy_(tensor)
state_dict.pop(key)
+ counter_after_loading_weights = time.perf_counter()
+ logger.info_once(
+ "Loading weights took %.2f seconds",
+ counter_after_loading_weights - counter_before_loading_weights,
+ scope="local",
+ )
if state_dict:
raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")
diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py
index ba708a098c0da..2021b68b8a60b 100644
--- a/vllm/model_executor/model_loader/utils.py
+++ b/vllm/model_executor/model_loader/utils.py
@@ -19,12 +19,6 @@ from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
-from vllm.model_executor.models.adapters import (
- as_embedding_model,
- as_reward_model,
- as_seq_cls_model,
- try_create_mm_pooling_model_cls,
-)
from vllm.model_executor.models.interfaces import SupportsQuant, supports_multimodal
from vllm.utils.platform_utils import is_pin_memory_available
@@ -88,6 +82,14 @@ def initialize_model(
def process_weights_after_loading(
model: nn.Module, model_config: ModelConfig, target_device: torch.device
) -> None:
+ if getattr(model, "process_weights_after_loading_already_called", False):
+ # In case `process_weights_after_loading` is called multiple times
+ # we'll skip it at later times
+ logger.debug_once(
+ "process_weights_after_loading already called for model %s", model
+ )
+ return
+
# to avoid circular dependency
from vllm.model_executor.model_loader.online_quantization import (
maybe_save_metadata_and_attributes_for_weight_reloading,
@@ -164,6 +166,13 @@ _MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]()
def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], str]:
+ from vllm.model_executor.models.adapters import (
+ as_embedding_model,
+ as_reward_model,
+ as_seq_cls_model,
+ try_create_mm_pooling_model_cls,
+ )
+
architectures = getattr(model_config.hf_config, "architectures", [])
model_cls, arch = model_config.registry.resolve_model_cls(
diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py
index 93986e5f2fc0a..4572ebe2ea11b 100644
--- a/vllm/model_executor/model_loader/weight_utils.py
+++ b/vllm/model_executor/model_loader/weight_utils.py
@@ -595,6 +595,9 @@ def safetensors_weights_iterator(
if safetensors_load_strategy == "eager":
loading_desc += " (eager)"
+ state_dict = {}
+ leftover_state_dict: dict[str, torch.Tensor] = {}
+
for st_file in tqdm(
hf_weights_files,
desc=loading_desc,
@@ -606,9 +609,11 @@ def safetensors_weights_iterator(
state_dict = load(f.read())
yield from state_dict.items()
elif safetensors_load_strategy == "torchao":
- if not torchao_version_at_least("0.14.0"):
+ # we can't load flattened torchao tensor subclasses directly into the model
+ # instead we reconstruct the subclasses here before returning
+ if not torchao_version_at_least("0.15.0"):
raise ValueError(
- "Please use torchao version >= 0.14.0 \
+ "Please use torchao version >= 0.15.0 \
to load torchao safetensors checkpoint"
)
from torchao.prototype.safetensors.safetensors_support import (
@@ -616,12 +621,20 @@ def safetensors_weights_iterator(
)
with safe_open(st_file, framework="pt") as f:
- state_dict = {}
for name in f.keys(): # noqa: SIM118
state_dict[name] = f.get_tensor(name)
+
+ # update with leftover tensor data from previous iteration, if any
+ state_dict.update(leftover_state_dict)
metadata = f.metadata()
- updated_state_dict = unflatten_tensor_state_dict(state_dict, metadata)
- yield from updated_state_dict.items()
+ # due to sharded checkpoints, we are not guaranteed that we have all
+ # tensor subclass data on one file
+ # state_dict has the leftover data from this step and we wait for
+ # missing information to be provided in a future iteration
+ unflattened_state_dict, leftover_state_dict = (
+ unflatten_tensor_state_dict(state_dict, metadata)
+ )
+ yield from unflattened_state_dict.items()
else:
with safe_open(st_file, framework="pt") as f:
for name in f.keys(): # noqa: SIM118
@@ -836,7 +849,11 @@ def gguf_quant_weights_iterator(
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""
Iterate over the quant weights in the model gguf files and convert
- them to torch tensors
+ them to torch tensors.
+ Be careful of the order of yielding weight types and weights data,
+ we have to yield all weight types first before yielding any weights.
+ Otherwise it would cause issue when loading weights with for packed
+ layer with different quant types.
"""
reader = gguf.GGUFReader(gguf_file)
@@ -846,7 +863,7 @@ def gguf_quant_weights_iterator(
weight_type = tensor.tensor_type
name = gguf_to_hf_name_map[tensor.name]
- if weight_type.name != "F32":
+ if weight_type.name not in ("F32", "BF16", "F16"):
weight_type_name = name.replace("weight", "qweight_type")
weight_type = torch.tensor(weight_type)
yield weight_type_name, weight_type
@@ -856,7 +873,7 @@ def gguf_quant_weights_iterator(
weight = tensor.data
weight_type = tensor.tensor_type
name = gguf_to_hf_name_map[tensor.name]
- if weight_type.name != "F32":
+ if weight_type.name not in ("F32", "BF16", "F16"):
name = name.replace("weight", "qweight")
param = torch.tensor(weight)
yield name, param
diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py
index 6f654f47495f7..4eb5665a71fc8 100644
--- a/vllm/model_executor/models/afmoe.py
+++ b/vllm/model_executor/models/afmoe.py
@@ -5,7 +5,6 @@
import typing
from collections.abc import Callable, Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -171,8 +170,6 @@ class AfmoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 131072,
head_dim: int | None = None,
rms_norm_eps: float = 1e-05,
@@ -202,7 +199,6 @@ class AfmoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
# Check if this is a local attention layer
@@ -246,8 +242,7 @@ class AfmoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config["rope_parameters"],
is_neox_style=True,
)
else:
@@ -303,14 +298,6 @@ class AfmoeDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
# DecoderLayers are created with `make_layers` which passes the prefix
@@ -323,8 +310,6 @@ class AfmoeDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
head_dim=config.head_dim,
rms_norm_eps=config.rms_norm_eps,
diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py
index 5872e8196eada..3d000f3ac3ab5 100644
--- a/vllm/model_executor/models/aimv2.py
+++ b/vllm/model_executor/models/aimv2.py
@@ -12,6 +12,7 @@ from vllm.attention.layer import MultiHeadAttention
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.distributed.utils import divide
from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
@@ -58,7 +59,7 @@ class AIMv2SwiGLUFFN(nn.Module):
class AIMv2PatchEmbed(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
- self.proj = nn.Conv2d(
+ self.proj = Conv2dLayer(
config.num_channels,
config.hidden_size,
kernel_size=(config.patch_size, config.patch_size),
diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py
index 0a8f21abb0a35..b75e91319bbad 100644
--- a/vllm/model_executor/models/apertus.py
+++ b/vllm/model_executor/models/apertus.py
@@ -27,7 +27,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -118,8 +117,6 @@ class ApertusAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -155,7 +152,6 @@ class ApertusAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -176,9 +172,7 @@ class ApertusAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
- self._init_rotary_emb(
- config, rope_scaling=rope_scaling, quant_config=quant_config
- )
+ self._init_rotary_emb(config, quant_config=quant_config)
sliding_window = None
if layer_types := getattr(config, "layer_types", None):
@@ -224,7 +218,6 @@ class ApertusAttention(nn.Module):
def _init_rotary_emb(
self,
config: ApertusConfig,
- rope_scaling: dict[str, Any] | None,
quant_config: QuantizationConfig | None,
) -> None:
is_neox_style = True
@@ -236,8 +229,7 @@ class ApertusAttention(nn.Module):
self.head_dim,
rotary_dim=int(self.partial_rotary_factor * self.head_dim),
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
partial_rotary_factor=self.partial_rotary_factor,
)
@@ -253,14 +245,6 @@ class ApertusDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -288,8 +272,6 @@ class ApertusDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py
index 20c3ff0754506..b3887b16f4d74 100644
--- a/vllm/model_executor/models/arcee.py
+++ b/vllm/model_executor/models/arcee.py
@@ -103,15 +103,6 @@ class ArceeDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Rotary embedding parameters (reuse LLaMA defaults)
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Determine if attention bias is needed (some variants use bias terms)
attention_bias = getattr(config, "attention_bias", False) or getattr(
@@ -133,8 +124,6 @@ class ArceeDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/arctic.py b/vllm/model_executor/models/arctic.py
index b5cc07a56535d..b75a254761d4e 100644
--- a/vllm/model_executor/models/arctic.py
+++ b/vllm/model_executor/models/arctic.py
@@ -292,7 +292,6 @@ class ArcticAttention(nn.Module):
self.kv_size = self.num_kv_heads * self.head_dim
self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
self.scaling = self.head_dim**-0.5
self.qkv_proj = QKVParallelLinear(
@@ -317,7 +316,7 @@ class ArcticAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=int(self.rope_theta),
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py
index 8991ef4c606b6..edf47270e5277 100644
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -136,7 +136,7 @@ class BaiChuanAttention(nn.Module):
hidden_size: int,
num_heads: int,
position_embedding: str,
- rope_theta: float = 10000,
+ rope_parameters: dict,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -150,7 +150,6 @@ class BaiChuanAttention(nn.Module):
self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
self.head_dim = hidden_size // self.total_num_heads
self.position_embedding = position_embedding
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
# pylint: disable=invalid-name
@@ -192,7 +191,7 @@ class BaiChuanAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=rope_parameters,
)
self.scaling = self.head_dim**-0.5
self.attn = Attention(
@@ -229,13 +228,12 @@ class BaiChuanDecoderLayer(nn.Module):
):
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.self_attn = BaiChuanAttention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
position_embedding=position_embedding,
- rope_theta=rope_theta,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py
index 024425bb24406..cc10e936a2d3d 100644
--- a/vllm/model_executor/models/bailing_moe.py
+++ b/vllm/model_executor/models/bailing_moe.py
@@ -135,9 +135,8 @@ class BailingAttention(nn.Module):
self.head_dim,
rotary_dim=self.rotary_dim,
max_position=config.max_position_embeddings,
- base=config.rope_theta,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
- rope_scaling=config.rope_scaling,
partial_rotary_factor=self.partial_rotary_factor,
)
diff --git a/vllm/model_executor/models/bamba.py b/vllm/model_executor/models/bamba.py
index e0a2defd5127e..4422bb5da98f4 100644
--- a/vllm/model_executor/models/bamba.py
+++ b/vllm/model_executor/models/bamba.py
@@ -138,8 +138,7 @@ class BambaMixerDecoderLayer(nn.Module):
else:
hidden_states, residual = self.input_layernorm(hidden_states, residual)
- output = torch.empty_like(hidden_states)
- self.mamba(hidden_states, output)
+ output = self.mamba(hidden_states)
# Fully Connected
hidden_states, residual = self.pre_ff_layernorm(output, residual)
hidden_states = self.feed_forward(hidden_states)
@@ -157,8 +156,6 @@ class BambaAttentionDecoderLayer(nn.Module):
prefix: str = "",
) -> None:
super().__init__()
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.hidden_size = config.hidden_size
tp_size = get_tensor_model_parallel_world_size()
@@ -179,7 +176,6 @@ class BambaAttentionDecoderLayer(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
if hasattr(config, "partial_rotary_factor"):
@@ -193,8 +189,7 @@ class BambaAttentionDecoderLayer(nn.Module):
head_size=self.head_dim,
rotary_dim=rotary_dim,
max_position=max_position_embeddings,
- rope_scaling=rope_scaling,
- base=rope_theta,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
dtype=torch.get_default_dtype(), # see impl of get_rope
)
diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py
index 2e4f73312efa3..f31f99c0592b2 100644
--- a/vllm/model_executor/models/blip.py
+++ b/vllm/model_executor/models/blip.py
@@ -12,6 +12,7 @@ from transformers import Blip2VisionConfig, BlipVisionConfig
from vllm.attention.layer import MultiHeadAttention
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -47,7 +48,7 @@ class BlipVisionEmbeddings(nn.Module):
self.class_embedding = nn.Parameter(torch.randn(1, 1, self.embed_dim))
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=3,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
diff --git a/vllm/model_executor/models/chameleon.py b/vllm/model_executor/models/chameleon.py
index fb7476c45fcdb..b5a6d00dc309f 100644
--- a/vllm/model_executor/models/chameleon.py
+++ b/vllm/model_executor/models/chameleon.py
@@ -22,6 +22,7 @@ from vllm.config.multimodal import BaseDummyOptions
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 SiluAndMul
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
@@ -264,8 +265,7 @@ class ChameleonAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any],
max_position_embeddings: int = 4096,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -292,7 +292,6 @@ class ChameleonAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -317,8 +316,7 @@ class ChameleonAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
@@ -368,14 +366,6 @@ class ChameleonDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 4096)
self.self_attn = ChameleonAttention(
@@ -384,8 +374,7 @@ class ChameleonDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=False,
@@ -438,14 +427,6 @@ class ChameleonSwinDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 4096)
self.self_attn = ChameleonAttention(
@@ -454,8 +435,7 @@ class ChameleonSwinDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=False,
@@ -549,7 +529,7 @@ class ChameleonVQVAEVectorQuantizer(nn.Module):
class ChameleonVQVAEEncoderConvDownsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
- self.conv = nn.Conv2d(
+ self.conv = Conv2dLayer(
in_channels, in_channels, kernel_size=3, stride=2, padding=0
)
@@ -577,23 +557,23 @@ class ChameleonVQVAEEncoderResnetBlock(nn.Module):
self.norm1 = torch.nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
- self.conv1 = torch.nn.Conv2d(
+ self.conv1 = Conv2dLayer(
in_channels, out_channels, kernel_size=3, stride=1, padding=1
)
self.norm2 = torch.nn.GroupNorm(
num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
)
self.dropout = torch.nn.Dropout(config.dropout)
- self.conv2 = torch.nn.Conv2d(
+ self.conv2 = Conv2dLayer(
out_channels, out_channels, kernel_size=3, stride=1, padding=1
)
if self.in_channels != self.out_channels:
if self.use_conv_shortcut:
- self.conv_shortcut = torch.nn.Conv2d(
+ self.conv_shortcut = Conv2dLayer(
in_channels, out_channels, kernel_size=3, stride=1, padding=1
)
else:
- self.nin_shortcut = torch.nn.Conv2d(
+ self.nin_shortcut = Conv2dLayer(
in_channels, out_channels, kernel_size=1, stride=1, padding=0
)
@@ -626,16 +606,16 @@ class ChameleonVQVAEEncoderAttnBlock(nn.Module):
self.norm = torch.nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
- self.q = torch.nn.Conv2d(
+ self.q = Conv2dLayer(
in_channels, in_channels, kernel_size=1, stride=1, padding=0
)
- self.k = torch.nn.Conv2d(
+ self.k = Conv2dLayer(
in_channels, in_channels, kernel_size=1, stride=1, padding=0
)
- self.v = torch.nn.Conv2d(
+ self.v = Conv2dLayer(
in_channels, in_channels, kernel_size=1, stride=1, padding=0
)
- self.proj_out = torch.nn.Conv2d(
+ self.proj_out = Conv2dLayer(
in_channels, in_channels, kernel_size=1, stride=1, padding=0
)
@@ -681,7 +661,7 @@ class ChameleonVQVAEEncoder(nn.Module):
latent_channels = config.latent_channels
channel_multiplier = config.channel_multiplier
- self.conv_in = torch.nn.Conv2d(
+ self.conv_in = Conv2dLayer(
in_channels, base_channels, kernel_size=3, stride=1, padding=1
)
@@ -738,7 +718,7 @@ class ChameleonVQVAEEncoder(nn.Module):
self.norm_out = torch.nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
- self.conv_out = torch.nn.Conv2d(
+ self.conv_out = Conv2dLayer(
block_in,
2 * latent_channels if double_latent else latent_channels,
kernel_size=3,
@@ -779,10 +759,8 @@ class ChameleonVQVAE(nn.Module):
super().__init__()
self.encoder = ChameleonVQVAEEncoder(config)
self.quantize = ChameleonVQVAEVectorQuantizer(config)
- self.quant_conv = torch.nn.Conv2d(config.latent_channels, config.embed_dim, 1)
- self.post_quant_conv = torch.nn.Conv2d(
- config.embed_dim, config.latent_channels, 1
- )
+ self.quant_conv = Conv2dLayer(config.latent_channels, config.embed_dim, 1)
+ self.post_quant_conv = Conv2dLayer(config.embed_dim, config.latent_channels, 1)
self.eval() # Chameleon's VQ model is frozen
def encode(
diff --git a/vllm/model_executor/models/chatglm.py b/vllm/model_executor/models/chatglm.py
index 5d6f5e9125a28..dbfcd62d0bcab 100644
--- a/vllm/model_executor/models/chatglm.py
+++ b/vllm/model_executor/models/chatglm.py
@@ -99,6 +99,7 @@ class GLMAttention(nn.Module):
# https://huggingface.co/zai-org/chatglm3-6b-32k/blob/e210410255278dd9d74463cf396ba559c0ef801c/modeling_chatglm.py#L141
rope_ratio = getattr(config, "rope_ratio", 1.0)
max_positions = getattr(config, "seq_length", 8192)
+ rope_parameters = {"rope_type": "default", "rope_theta": 10000 * rope_ratio}
# NOTE: zai-org/cogagent-9b-20241220 uses original_rope=False,
# which is equivalent to is_neox_style=True
is_neox_style = not config.original_rope
@@ -106,7 +107,7 @@ class GLMAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim // 2,
max_position=max_positions,
- base=10000 * rope_ratio,
+ rope_parameters=rope_parameters,
is_neox_style=is_neox_style,
)
self.attn = Attention(
diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py
index 77bb178519813..5ed920927c772 100644
--- a/vllm/model_executor/models/commandr.py
+++ b/vllm/model_executor/models/commandr.py
@@ -156,8 +156,6 @@ class CohereAttention(nn.Module):
self.max_position_embeddings = getattr(
config, "model_max_length", None
) or getattr(config, "max_position_embeddings", 8192)
- self.rope_theta = config.rope_theta
- self.rope_scaling = getattr(config, "rope_scaling", None)
self.use_qk_norm = getattr(config, "use_qk_norm", False)
self.qkv_proj = QKVParallelLinear(
self.hidden_size,
@@ -179,8 +177,7 @@ class CohereAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=self.rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=False,
)
diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py
index 66b246878b0aa..3cf4bf991e667 100644
--- a/vllm/model_executor/models/config.py
+++ b/vllm/model_executor/models/config.py
@@ -8,6 +8,7 @@ import vllm.envs as envs
from vllm.logger import init_logger
from vllm.model_executor.models import ModelRegistry
from vllm.platforms import current_platform
+from vllm.transformers_utils.config import set_default_rope_theta
from vllm.utils.math_utils import cdiv, round_up
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentionSpec
@@ -46,8 +47,7 @@ class GteNewModelConfig(VerifyAndUpdateConfig):
"head_size": head_dim,
"rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
"max_position": config.max_position_embeddings,
- "base": config.rope_theta,
- "rope_scaling": getattr(config, "rope_scaling", None),
+ "rope_parameters": config.rope_parameters,
}
@@ -78,12 +78,13 @@ class JinaRobertaModelConfig(VerifyAndUpdateConfig):
if not model_config.enforce_eager:
max_position = round_up(max_position, 8)
+ set_default_rope_theta(config, default_theta=config.rotary_emb_base)
+
config.rotary_kwargs = {
"head_size": head_dim,
"rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
"max_position": max_position,
- "base": getattr(config, "rope_theta", config.rotary_emb_base),
- "rope_scaling": getattr(config, "rope_scaling", None),
+ "rope_parameters": config.rope_parameters,
}
@@ -117,18 +118,20 @@ class NomicBertModelConfig(VerifyAndUpdateConfig):
head_dim = config.hidden_size // config.num_attention_heads
rotary_emb_dim = int(head_dim * config.rotary_emb_fraction)
max_trained_positions = getattr(config, "max_trained_positions", 2048)
+
+ set_default_rope_theta(config, default_theta=config.rotary_emb_base)
+
config.rotary_kwargs = {
"head_size": head_dim,
"rotary_dim": rotary_emb_dim,
"max_position": max_trained_positions,
- "base": getattr(config, "rope_theta", config.rotary_emb_base),
- "rope_scaling": getattr(config, "rope_scaling", None),
+ "rope_parameters": config.rope_parameters,
}
# we ignore config.rotary_scaling_factor so that for datasets shorter
# than max_trained_positions 2048, the results are consistent
# with SentenceTransformer.
- # The context extension uses vllm style rope_theta and rope_scaling.
+ # The context extension uses vllm style rope_theta and rope_parameters.
# See #17785 #18755
if (
not vllm_config.model_config.hf_overrides
@@ -172,7 +175,7 @@ class NomicBertModelConfig(VerifyAndUpdateConfig):
if hasattr(hf_text_config, "max_model_len"):
delattr(hf_text_config, "max_model_len")
hf_text_config.max_position_embeddings = max_trained_positions
- hf_text_config.rope_scaling = config.rotary_kwargs["rope_scaling"]
+ hf_text_config.rope_parameters = config.rotary_kwargs["rope_parameters"]
# The priority of sentence_bert_config.json is higher
# than max_position_embeddings
@@ -246,8 +249,7 @@ class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
"head_size": head_dim,
"rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
"max_position": config.max_position_embeddings,
- "base": config.rope_theta,
- "rope_scaling": getattr(config, "rope_scaling", None),
+ "rope_parameters": config.rope_parameters,
}
diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py
index 528ef4f76742d..2c729019081a4 100644
--- a/vllm/model_executor/models/dbrx.py
+++ b/vllm/model_executor/models/dbrx.py
@@ -197,7 +197,10 @@ class DbrxAttention(nn.Module):
self.head_dim = self.d_model // self.total_num_heads
self.total_num_kv_heads = config.attn_config.kv_n_heads
self.clip_qkv = config.attn_config.clip_qkv
- self.rope_theta = config.attn_config.rope_theta
+ rope_parameters = {
+ "rope_type": "default",
+ "rope_theta": int(config.attn_config.rope_theta),
+ }
self.max_position = config.max_seq_len
# pylint: disable=invalid-name
@@ -221,7 +224,7 @@ class DbrxAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position,
- base=int(self.rope_theta),
+ rope_parameters=rope_parameters,
is_neox_style=True,
)
diff --git a/vllm/model_executor/models/deepencoder.py b/vllm/model_executor/models/deepencoder.py
index e62a57eccc953..8f1660891fcbf 100644
--- a/vllm/model_executor/models/deepencoder.py
+++ b/vllm/model_executor/models/deepencoder.py
@@ -19,6 +19,7 @@ import torch.nn.functional as F
from transformers import CLIPVisionConfig
from vllm.attention.layer import MultiHeadAttention
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -133,14 +134,14 @@ class ImageEncoderViT(nn.Module):
self.blocks.append(block)
self.neck = nn.Sequential(
- nn.Conv2d(
+ Conv2dLayer(
embed_dim,
out_chans,
kernel_size=1,
bias=False,
),
LayerNorm2d(out_chans),
- nn.Conv2d(
+ Conv2dLayer(
out_chans,
out_chans,
kernel_size=3,
@@ -150,8 +151,10 @@ class ImageEncoderViT(nn.Module):
LayerNorm2d(out_chans),
)
- self.net_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False)
- self.net_3 = nn.Conv2d(
+ self.net_2 = Conv2dLayer(
+ 256, 512, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.net_3 = Conv2dLayer(
512, 1024, kernel_size=3, stride=2, padding=1, bias=False
)
@@ -500,7 +503,7 @@ class PatchEmbed(nn.Module):
"""
super().__init__()
- self.proj = nn.Conv2d(
+ self.proj = Conv2dLayer(
in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
)
diff --git a/vllm/model_executor/models/deepseek_eagle.py b/vllm/model_executor/models/deepseek_eagle.py
index 3fb04c3b70dd1..4d7a37292cb02 100644
--- a/vllm/model_executor/models/deepseek_eagle.py
+++ b/vllm/model_executor/models/deepseek_eagle.py
@@ -8,7 +8,6 @@ import torch.nn as nn
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
-from vllm.distributed.parallel_state import get_pp_group
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -172,10 +171,6 @@ class DeepseekV2Model(nn.Module):
)
break
else:
- # if PP disabled then draft will share embed with target
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
- continue
-
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index e8ee9951d6119..7cfd381592b49 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/vllm/model_executor/models/deepseek_v2.py
@@ -27,7 +27,6 @@
import typing
from collections.abc import Callable, Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -111,8 +110,6 @@ class DeepseekAttention(nn.Module):
config: DeepseekV2Config | DeepseekV3Config,
hidden_size: int,
num_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -139,7 +136,6 @@ class DeepseekAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -162,8 +158,7 @@ class DeepseekAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -292,7 +287,10 @@ class DeepseekV2MoE(nn.Module):
)
self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
- if config.n_shared_experts is None or self.is_rocm_aiter_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
@@ -332,7 +330,7 @@ class DeepseekV2MoE(nn.Module):
num_redundant_experts=self.n_redundant_experts,
is_sequence_parallel=self.is_sequence_parallel,
n_shared_experts=config.n_shared_experts
- if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
+ if self.is_fusion_moe_shared_experts_enabled
else None,
)
@@ -409,8 +407,6 @@ class DeepseekV2Attention(nn.Module):
v_head_dim: int,
q_lora_rank: int,
kv_lora_rank: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -430,7 +426,6 @@ class DeepseekV2Attention(nn.Module):
assert num_heads % tp_size == 0
self.num_local_heads = num_heads // tp_size
self.scaling = self.qk_head_dim**-0.5
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
assert topk_indices_buffer is None, (
"topk_indices_buffer is not \
@@ -485,21 +480,20 @@ class DeepseekV2Attention(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
)
- if rope_scaling:
- rope_scaling["rope_type"] = "deepseek_yarn"
+ if config.rope_parameters["rope_type"] != "default":
+ config.rope_parameters["rope_type"] = "deepseek_yarn"
self.rotary_emb = get_rope(
qk_rope_head_dim,
rotary_dim=qk_rope_head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=False,
)
- if rope_scaling:
- mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
- scaling_factor = rope_scaling["factor"]
+ if config.rope_parameters["rope_type"] != "default":
+ 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))
self.scaling = self.scaling * mscale * mscale
@@ -600,6 +594,7 @@ def sparse_attn_indexer(
) -> torch.Tensor:
# careful! this will be None in dummy run
attn_metadata = get_forward_context().attn_metadata
+ fp8_dtype = current_platform.fp8_dtype()
# assert isinstance(attn_metadata, dict)
if not isinstance(attn_metadata, dict):
return sparse_attn_indexer_fake(
@@ -639,7 +634,7 @@ def sparse_attn_indexer(
k_fp8 = torch.empty(
[chunk.total_seq_lens, head_dim],
device=k.device,
- dtype=torch.float8_e4m3fn,
+ dtype=fp8_dtype,
)
k_scale = torch.empty(
[chunk.total_seq_lens, 4],
@@ -653,7 +648,12 @@ def sparse_attn_indexer(
chunk.block_table,
chunk.cu_seq_lens,
)
- logits = fp8_mqa_logits(
+ fp8_mqa_logits_func = fp8_mqa_logits
+ if current_platform.is_rocm():
+ from vllm.attention.ops.rocm_aiter_mla_sparse import rocm_fp8_mqa_logits
+
+ fp8_mqa_logits_func = rocm_fp8_mqa_logits
+ logits = fp8_mqa_logits_func(
q_fp8[chunk.token_start : chunk.token_end],
(k_fp8, k_scale.view(torch.float32)),
weights[chunk.token_start : chunk.token_end],
@@ -698,7 +698,14 @@ def sparse_attn_indexer(
next_n = padded_q_fp8_decode_tokens.shape[1]
assert batch_size == decode_metadata.seq_lens.shape[0]
num_padded_tokens = batch_size * next_n
- logits = fp8_paged_mqa_logits(
+ fp8_paged_mqa_logits_func = fp8_paged_mqa_logits
+ if current_platform.is_rocm():
+ from vllm.attention.ops.rocm_aiter_mla_sparse import (
+ rocm_fp8_paged_mqa_logits,
+ )
+
+ fp8_paged_mqa_logits_func = rocm_fp8_paged_mqa_logits
+ logits = fp8_paged_mqa_logits_func(
padded_q_fp8_decode_tokens,
kv_cache,
weights[:num_padded_tokens],
@@ -755,7 +762,8 @@ def sparse_attn_indexer_fake(
_flattened_kv = torch.empty(
[total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
)
- _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous()
+ fp8_dtype = current_platform.fp8_dtype()
+ _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
_k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()
return topk_indices_buffer
@@ -846,8 +854,8 @@ class Indexer(nn.Module):
)
q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
- q = torch.cat([q_pe, q_nope], dim=-1)
- k = torch.cat([k_pe.squeeze(1), k_nope], dim=-1)
+ q = torch.cat([q_pe.squeeze(0), q_nope], dim=-1)
+ k = torch.cat([k_pe.squeeze((0, 2)), k_nope], dim=-1)
# we only quant q here since k quant is fused with cache insertion
q = q.view(-1, self.head_dim)
@@ -903,8 +911,6 @@ class DeepseekV2MLAAttention(nn.Module):
v_head_dim: int,
q_lora_rank: int | None,
kv_lora_rank: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -927,7 +933,6 @@ class DeepseekV2MLAAttention(nn.Module):
self.num_local_heads = num_heads // tp_size
self.scaling = self.qk_head_dim**-0.5
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
if self.q_lora_rank is not None:
@@ -981,25 +986,31 @@ class DeepseekV2MLAAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
- if rope_scaling:
- rope_scaling["rope_type"] = "deepseek_yarn"
+ if config.rope_parameters["rope_type"] != "default":
+ config.rope_parameters["rope_type"] = "deepseek_yarn"
self.rotary_emb = get_rope(
qk_rope_head_dim,
rotary_dim=qk_rope_head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=False,
)
- if rope_scaling:
- mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
- scaling_factor = rope_scaling["factor"]
+ if config.rope_parameters["rope_type"] != "default":
+ 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))
self.scaling = self.scaling * mscale * mscale
self.is_v32 = hasattr(config, "index_topk")
if self.is_v32:
+ self.indexer_rope_emb = get_rope(
+ qk_rope_head_dim,
+ rotary_dim=qk_rope_head_dim,
+ max_position=max_position_embeddings,
+ rope_parameters=config.rope_parameters,
+ is_neox_style=True,
+ )
self.indexer = Indexer(
vllm_config,
config,
@@ -1011,6 +1022,7 @@ class DeepseekV2MLAAttention(nn.Module):
f"{prefix}.indexer",
)
else:
+ self.indexer_rope_emb = None
self.indexer = None
mla_modules = MLAModules(
@@ -1028,6 +1040,7 @@ class DeepseekV2MLAAttention(nn.Module):
q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
q_proj=self.q_proj if self.q_lora_rank is None else None,
indexer=self.indexer,
+ indexer_rotary_emb=self.indexer_rope_emb,
is_sparse=self.is_v32,
topk_indices_buffer=topk_indices_buffer,
)
@@ -1073,8 +1086,6 @@ class DeepseekV2DecoderLayer(nn.Module):
parallel_config = vllm_config.parallel_config
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
moe_layer_freq = getattr(config, "moe_layer_freq", 1)
# DecoderLayers are created with `make_layers` which passes the prefix
@@ -1107,8 +1118,6 @@ class DeepseekV2DecoderLayer(nn.Module):
v_head_dim=v_head_dim,
q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,
kv_lora_rank=kv_lora_rank,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py
index d24da0c42a254..e65c275106a4e 100644
--- a/vllm/model_executor/models/dots1.py
+++ b/vllm/model_executor/models/dots1.py
@@ -27,7 +27,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -202,8 +201,6 @@ class Dots1Attention(nn.Module):
num_heads: int,
num_kv_heads: int,
config: Dots1Config,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -229,7 +226,6 @@ class Dots1Attention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
attention_bias = config.attention_bias
@@ -255,8 +251,7 @@ class Dots1Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -296,8 +291,6 @@ class Dots1DecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
layer_idx = int(prefix.split(sep=".")[-1])
self.layer_idx = layer_idx
@@ -307,8 +300,6 @@ class Dots1DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
config=config,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/dots_ocr.py b/vllm/model_executor/models/dots_ocr.py
index f46caaa095c6a..5460018d0d67a 100644
--- a/vllm/model_executor/models/dots_ocr.py
+++ b/vllm/model_executor/models/dots_ocr.py
@@ -22,6 +22,7 @@ from vllm.distributed.parallel_state import (
get_tensor_model_parallel_world_size,
)
from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
@@ -305,7 +306,6 @@ class DotsVisionAttention(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -323,7 +323,6 @@ class DotsVisionAttention(nn.Module):
rotary_pos_emb: torch.Tensor | None = None,
*,
max_seqlen: int | None = None,
- seqlens: list[int] | None = None,
) -> torch.Tensor:
# [S, C] -> [S, B=1, C]
x = hidden_states.unsqueeze(1)
@@ -373,16 +372,6 @@ class DotsVisionAttention(nn.Module):
out_i = out_i.permute(0, 2, 1, 3)
outputs.append(out_i)
context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
-
- attn_bias = BlockDiagonalMask.from_seqlens(
- q_seqlen=seqlens, kv_seqlen=None, device=q.device
- )
- context_layer = xops.memory_efficient_attention_forward(
- q, k, v, attn_bias=attn_bias, p=0, scale=None
- )
else:
raise RuntimeError("Unsupported attention backend")
@@ -471,7 +460,7 @@ class DotsPatchEmbed(nn.Module):
self.temporal_patch_size = config.temporal_patch_size
self.embed_dim = config.embed_dim
self.config = config
- self.proj = nn.Conv2d(
+ self.proj = Conv2dLayer(
config.num_channels,
config.embed_dim,
kernel_size=(config.patch_size, config.patch_size),
@@ -544,14 +533,12 @@ class DotsVisionBlock(nn.Module):
cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor,
max_seqlen: int | None = None,
- seqlens: list[int] | None = None,
) -> torch.Tensor:
hidden_states = hidden_states + self.attn(
self.norm1(hidden_states),
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
return hidden_states
@@ -662,18 +649,14 @@ class DotsVisionTransformer(nn.Module):
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
return rotary_pos_emb
- def compute_attn_mask_seqlen(
- self, cu_seqlens: torch.Tensor
- ) -> tuple[int | None, list[int] | None]:
- max_seqlen, seqlens = None, None
+ def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
+ max_seqlen = None
if (
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
):
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
- return max_seqlen, seqlens
+ return max_seqlen
def forward(
self, hidden_states: torch.Tensor, grid_thw: list[list[int]]
@@ -693,14 +676,13 @@ class DotsVisionTransformer(nn.Module):
)
cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
- max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
for blk in self.blocks:
hidden_states = blk(
hidden_states,
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
if self.post_trunk_norm is not None:
diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py
index f2999968669f6..a7df3509e3ecd 100644
--- a/vllm/model_executor/models/ernie45_moe.py
+++ b/vllm/model_executor/models/ernie45_moe.py
@@ -62,6 +62,7 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
from .utils import (
@@ -232,9 +233,8 @@ class Ernie4_5_MoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
+ rope_parameters: dict[str, Any],
head_dim: int | None = None,
- rope_theta: float = 500000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 131072,
rms_norm_eps: float = 1e-05,
qkv_bias: bool = False,
@@ -266,7 +266,6 @@ class Ernie4_5_MoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -291,9 +290,8 @@ class Ernie4_5_MoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
+ rope_parameters=rope_parameters,
is_neox_style=False,
- rope_scaling=rope_scaling,
)
self.attn = Attention(
self.num_heads,
@@ -333,16 +331,14 @@ class Ernie4_5_MoeDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 500000)
- rope_scaling = getattr(config, "rope_scaling", None)
+ set_default_rope_theta(config, default_theta=500000)
max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
self.self_attn = Ernie4_5_MoeAttention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
head_dim=getattr(config, "head_dim", None),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
rms_norm_eps=config.rms_norm_eps,
qkv_bias=getattr(config, "use_bias", False),
diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py
index daa5bf03ea4a9..07b34fbc8addb 100644
--- a/vllm/model_executor/models/ernie45_vl.py
+++ b/vllm/model_executor/models/ernie45_vl.py
@@ -214,7 +214,6 @@ class Ernie4_5_VisionAttention(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -259,7 +258,6 @@ class Ernie4_5_VisionAttention(nn.Module):
cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor,
max_seqlen: int | None = None, # Only used for Flash Attention
- seqlens: list[int] | None = None, # Only used for xFormers
) -> torch.Tensor:
# [s, b, c] --> [s, b, head * 3 * head_dim]
x, _ = self.qkv(x)
@@ -311,20 +309,6 @@ class Ernie4_5_VisionAttention(nn.Module):
context_layer = rearrange(
context_layer, "b s h d -> s b (h d)"
).contiguous()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
-
- attn_bias = BlockDiagonalMask.from_seqlens(
- q_seqlen=seqlens, kv_seqlen=None, device=q.device
- )
-
- context_layer = xops.memory_efficient_attention_forward(
- q, k, v, attn_bias=attn_bias, p=0, scale=None
- )
- context_layer = rearrange(
- context_layer, "b s h d -> s b (h d)"
- ).contiguous()
output, _ = self.proj(context_layer)
return output
@@ -404,14 +388,12 @@ class Ernie4_5_VisionBlock(nn.Module):
cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor,
max_seqlen: int | None = None, # Only used for Flash Attention
- seqlens: list[int] | None = None, # Only used for xFormers
) -> torch.Tensor:
hidden_states = hidden_states + self.attn(
self.norm1(hidden_states),
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
return hidden_states
@@ -562,18 +544,14 @@ class Ernie4_5_VisionTransformer(nn.Module):
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
return rotary_pos_emb
- def compute_attn_mask_seqlen(
- self, cu_seqlens: torch.Tensor
- ) -> tuple[int | None, list[int] | None]:
- max_seqlen, seqlens = None, None
+ def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
+ max_seqlen = None
if (
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
):
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
- return max_seqlen, seqlens
+ return max_seqlen
def forward(
self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, num_pad=0
@@ -598,8 +576,8 @@ class Ernie4_5_VisionTransformer(nn.Module):
if hidden_states.ndim == 2:
hidden_states = hidden_states.unsqueeze(dim=1)
- # pre-compute seqlens for attn mask to reduce cuMemcpy operations
- max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+ # pre-compute max_seqlen for attn mask to reduce cuMemcpy operations
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
for i, blk in enumerate(self.blocks):
hidden_states = blk(
@@ -607,7 +585,6 @@ class Ernie4_5_VisionTransformer(nn.Module):
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
final_output = self.ln(hidden_states)
diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py
index e8ef86f9b7f01..50e033d77606d 100644
--- a/vllm/model_executor/models/ernie45_vl_moe.py
+++ b/vllm/model_executor/models/ernie45_vl_moe.py
@@ -58,6 +58,7 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
from .ernie45_moe import Ernie4_5_MoeMLP
from .interfaces import SupportsPP
@@ -91,9 +92,8 @@ class Ernie4_5_VLMoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
+ rope_parameters: dict[str, Any],
head_dim: int | None = None,
- rope_theta: float = 500000,
- rope_scaling: dict[str, Any] | None = None,
freq_allocation: int = 20,
max_position_embeddings: int = 131072,
rms_norm_eps: float = 1e-05,
@@ -126,7 +126,6 @@ class Ernie4_5_VLMoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -155,7 +154,7 @@ class Ernie4_5_VLMoeAttention(nn.Module):
head_size=self.head_dim,
rotary_dim=self.head_dim,
max_position_embeddings=max_position_embeddings,
- base=rope_theta,
+ base=rope_parameters["rope_theta"],
is_neox_style=False,
dtype=torch.get_default_dtype(),
mrope_section=[h_rope, w_rope, t_rope],
@@ -413,8 +412,7 @@ class Ernie4_5_VLMoeDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 500000)
- rope_scaling = getattr(config, "rope_scaling", None)
+ set_default_rope_theta(config, default_theta=500000)
freq_allocation = getattr(config, "freq_allocation", 20)
max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
@@ -423,8 +421,7 @@ class Ernie4_5_VLMoeDecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
head_dim=getattr(config, "head_dim", None),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
freq_allocation=freq_allocation,
max_position_embeddings=max_position_embeddings,
rms_norm_eps=config.rms_norm_eps,
diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py
index 6c56bfc433c7a..d13275488fe99 100644
--- a/vllm/model_executor/models/exaone.py
+++ b/vllm/model_executor/models/exaone.py
@@ -27,7 +27,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -113,8 +112,6 @@ class ExaoneAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -144,7 +141,6 @@ class ExaoneAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -173,8 +169,7 @@ class ExaoneAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
)
self.attn = Attention(
@@ -207,8 +202,6 @@ class ExaoneBlockAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -221,8 +214,6 @@ class ExaoneBlockAttention(nn.Module):
hidden_size=hidden_size,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=bias,
@@ -251,14 +242,6 @@ class ExaoneDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -272,8 +255,6 @@ class ExaoneDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py
index b89e168ada20e..70f3cce2b7c56 100644
--- a/vllm/model_executor/models/exaone4.py
+++ b/vllm/model_executor/models/exaone4.py
@@ -23,7 +23,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -52,6 +51,7 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
from .interfaces import SupportsLoRA, SupportsPP
from .utils import (
@@ -110,8 +110,6 @@ class Exaone4Attention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 1000000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -141,7 +139,6 @@ class Exaone4Attention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -176,12 +173,12 @@ class Exaone4Attention(nn.Module):
# apply rotary embeddings to every layer in full attention models
self.apply_rope_all_layers = "sliding_attention" not in config.layer_types
+ set_default_rope_theta(config, default_theta=1000000)
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
)
self.attn = Attention(
@@ -227,14 +224,6 @@ class Exaone4DecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 1000000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -249,8 +238,6 @@ class Exaone4DecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/falcon.py b/vllm/model_executor/models/falcon.py
index 85acdff3d96b4..dc2d51f340c8c 100644
--- a/vllm/model_executor/models/falcon.py
+++ b/vllm/model_executor/models/falcon.py
@@ -164,13 +164,12 @@ class FalconAttention(nn.Module):
)
if self.use_rotary:
- rope_theta = getattr(config, "rope_theta", 10000)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py
index 3653425b8e1ca..9433f0d1b4a49 100644
--- a/vllm/model_executor/models/falcon_h1.py
+++ b/vllm/model_executor/models/falcon_h1.py
@@ -35,6 +35,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
from .interfaces import (
HasInnerState,
@@ -198,10 +199,8 @@ class FalconH1SSMDecoderLayer(nn.Module):
residual: torch.Tensor | None,
**kwargs,
):
- output = torch.empty_like(hidden_states)
- self.mamba(
+ output = self.mamba(
hidden_states,
- output,
mup_vector=self.mup_vector,
)
return output, residual
@@ -216,8 +215,7 @@ class FalconH1AttentionDecoderLayer(nn.Module):
prefix: str = "",
) -> None:
super().__init__()
- rope_theta = getattr(config, "rope_theta", 1e11)
- rope_scaling = getattr(config, "rope_scaling", None)
+ set_default_rope_theta(config, default_theta=1e11)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.hidden_size = config.hidden_size
tp_size = get_tensor_model_parallel_world_size()
@@ -242,7 +240,6 @@ class FalconH1AttentionDecoderLayer(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
if hasattr(config, "partial_rotary_factor"):
@@ -256,8 +253,7 @@ class FalconH1AttentionDecoderLayer(nn.Module):
head_size=self.head_dim,
rotary_dim=rotary_dim,
max_position=max_position_embeddings,
- rope_scaling=rope_scaling,
- base=rope_theta,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
dtype=None, # see impl of get_rope
)
diff --git a/vllm/model_executor/models/gemma.py b/vllm/model_executor/models/gemma.py
index 7aaae7c503b58..00c7f59a08094 100644
--- a/vllm/model_executor/models/gemma.py
+++ b/vllm/model_executor/models/gemma.py
@@ -20,6 +20,7 @@
from collections.abc import Iterable
from functools import cache
from itertools import islice
+from typing import Any
import torch
from torch import nn
@@ -127,8 +128,8 @@ class GemmaAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
head_dim: int,
+ rope_parameters: dict[str, Any],
max_position_embeddings: int = 8192,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
@@ -153,7 +154,6 @@ class GemmaAttention(nn.Module):
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.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -176,7 +176,7 @@ class GemmaAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
@@ -218,7 +218,7 @@ class GemmaDecoderLayer(nn.Module):
num_kv_heads=config.num_key_value_heads,
head_dim=config.head_dim,
max_position_embeddings=config.max_position_embeddings,
- rope_theta=config.rope_theta,
+ rope_parameters=config.rope_parameters,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
diff --git a/vllm/model_executor/models/gemma2.py b/vllm/model_executor/models/gemma2.py
index 4d5d6cbb37c62..9b6cfe6932300 100644
--- a/vllm/model_executor/models/gemma2.py
+++ b/vllm/model_executor/models/gemma2.py
@@ -107,7 +107,6 @@ class Gemma2Attention(nn.Module):
num_kv_heads: int,
head_dim: int,
max_position_embeddings: int,
- rope_theta: float,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
attn_logits_soft_cap: float | None = None,
@@ -134,7 +133,6 @@ class Gemma2Attention(nn.Module):
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = config.query_pre_attn_scalar**-0.5
- self.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -156,7 +154,7 @@ class Gemma2Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
@@ -206,7 +204,6 @@ class Gemma2DecoderLayer(nn.Module):
num_kv_heads=config.num_key_value_heads,
head_dim=config.head_dim,
max_position_embeddings=config.max_position_embeddings,
- rope_theta=config.rope_theta,
cache_config=cache_config,
quant_config=quant_config,
attn_logits_soft_cap=config.attn_logit_softcapping,
diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py
index 357e61a4e78bf..4ad6fc89dcaf2 100644
--- a/vllm/model_executor/models/gemma3.py
+++ b/vllm/model_executor/models/gemma3.py
@@ -155,25 +155,30 @@ class Gemma3Attention(nn.Module):
self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
layer_idx = extract_layer_index(prefix)
- self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
+ layer_type = config.layer_types[layer_idx]
+ self.is_sliding = layer_type == "sliding_attention"
sliding_window = config.sliding_window if self.is_sliding else None
# Initialize the rotary embedding.
- if self.is_sliding:
- # Local attention. Override the values in config.json.
- self.rope_theta = config.rope_local_base_freq
- self.rope_scaling = {"rope_type": "default"}
+ if layer_type in config.rope_parameters:
+ # Transformers v5 rope config.
+ rope_parameters = config.rope_parameters[layer_type]
else:
+ # Transformers v4 rope config.
# Global attention. Use the values in config.json.
- self.rope_theta = config.rope_theta
- self.rope_scaling = config.rope_scaling
+ rope_parameters = config.rope_parameters
+ # Local attention. Override the values in config.json.
+ if self.is_sliding:
+ rope_parameters = dict(
+ rope_type="default", rope_theta=config.rope_local_base_freq
+ )
+
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=rope_parameters,
is_neox_style=True,
- rope_scaling=self.rope_scaling,
)
if getattr(config, "is_causal", True):
diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py
index 8e2bbe8f7990c..43c69e5e13992 100644
--- a/vllm/model_executor/models/gemma3_mm.py
+++ b/vllm/model_executor/models/gemma3_mm.py
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from collections.abc import Iterable, Mapping, Sequence
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal
import torch
from torch import nn
@@ -20,12 +20,7 @@ from vllm.multimodal.inputs import (
MultiModalFieldConfig,
MultiModalKwargsItems,
)
-from vllm.multimodal.parse import (
- ImageEmbeddingItems,
- ImageProcessorItems,
- ImageSize,
- MultiModalDataItems,
-)
+from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
from vllm.multimodal.processing import (
BaseMultiModalProcessor,
BaseProcessingInfo,
@@ -76,15 +71,7 @@ class Gemma3ImagePixelInputs(TensorSchema):
num_patches: Annotated[torch.Tensor, TensorShape("bn")]
-class Gemma3ImageEmbeddingInputs(TensorSchema):
- type: Literal["image_embeds"] = "image_embeds"
- image_embeds: Annotated[
- torch.Tensor,
- TensorShape("ni", "nf", "hs"),
- ]
-
-
-Gemma3ImageInputs: TypeAlias = Gemma3ImagePixelInputs | Gemma3ImageEmbeddingInputs
+Gemma3ImageInputs = Gemma3ImagePixelInputs
class Gemma3ProcessingInfo(BaseProcessingInfo):
@@ -191,9 +178,8 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):
def get_image_repl(
self,
*,
- image_width: int | None,
- image_height: int | None,
- num_crops: int | None = None,
+ image_width: int,
+ image_height: int,
processor: Gemma3Processor | None,
) -> PromptUpdateDetails[str]:
if processor is None:
@@ -201,13 +187,11 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):
boi_token = processor.boi_token
- if num_crops is None:
- assert image_width is not None and image_height is not None
- num_crops = self.get_num_crops(
- image_width=image_width,
- image_height=image_height,
- processor=processor,
- )
+ num_crops = self.get_num_crops(
+ image_width=image_width,
+ image_height=image_height,
+ processor=processor,
+ )
if num_crops == 0:
image_text = boi_token
@@ -337,7 +321,6 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
return dict(
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
num_patches=MultiModalFieldConfig.batched("image"),
- image_embeds=MultiModalFieldConfig.batched("image"),
)
def _get_prompt_updates(
@@ -350,19 +333,7 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
image_token = hf_processor.boi_token
def get_replacement_gemma3(item_idx: int):
- images = mm_items.get_items(
- "image", (ImageEmbeddingItems, ImageProcessorItems)
- )
-
- if isinstance(images, ImageEmbeddingItems):
- # For image embedding inputs, only support no crops cases
- # since it's not supported in hf processor anyway
- return self.info.get_image_repl(
- image_width=None,
- image_height=None,
- num_crops=0,
- processor=hf_processor,
- )
+ images = mm_items.get_items("image", ImageProcessorItems)
image_size = images.get_image_size(item_idx)
return self.info.get_image_repl(
@@ -586,19 +557,17 @@ class Gemma3ForConditionalGeneration(
pixel_values = kwargs.pop("pixel_values", None)
num_patches = kwargs.pop("num_patches", None)
image_embeds = kwargs.pop("image_embeds", None)
+ assert image_embeds is None, "Gemma3 does not support image_embeds."
+ if pixel_values is None:
+ return None
- if pixel_values is not None:
- image_size = self.config.vision_config.image_size
- return Gemma3ImagePixelInputs(
- pixel_values=pixel_values,
- num_patches=num_patches,
- resolve_bindings={"h": image_size, "w": image_size},
- )
- elif image_embeds is not None:
- return Gemma3ImageEmbeddingInputs(
- image_embeds=image_embeds,
- type="image_embeds",
- )
+ image_size = self.config.vision_config.image_size
+
+ return Gemma3ImagePixelInputs(
+ pixel_values=pixel_values,
+ num_patches=num_patches,
+ resolve_bindings={"h": image_size, "w": image_size},
+ )
def _image_pixels_to_features(
self,
@@ -610,9 +579,7 @@ class Gemma3ForConditionalGeneration(
def _process_image_input(
self,
image_input: Gemma3ImageInputs,
- ) -> torch.Tensor | list[torch.Tensor]:
- if image_input["type"] == "image_embeds":
- return image_input["image_embeds"]
+ ) -> list[torch.Tensor]:
assert self.vision_tower is not None
pixel_values = image_input["pixel_values"]
@@ -636,6 +603,26 @@ class Gemma3ForConditionalGeneration(
return self._process_image_input(image_input)
+ def embed_input_ids(
+ self,
+ input_ids: torch.Tensor,
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
+ *,
+ is_multimodal: torch.Tensor | None = None,
+ handle_oov_mm_token: bool = True,
+ ) -> torch.Tensor:
+ # Early return for text-only inference (no multimodal data)
+ if multimodal_embeddings is None or is_multimodal is None:
+ return super().embed_input_ids(input_ids)
+
+ # Use interface default with OOV handling enabled
+ return super().embed_input_ids(
+ input_ids,
+ multimodal_embeddings=multimodal_embeddings,
+ is_multimodal=is_multimodal,
+ handle_oov_mm_token=handle_oov_mm_token,
+ )
+
def forward(
self,
input_ids: torch.Tensor,
@@ -657,69 +644,6 @@ class Gemma3ForConditionalGeneration(
return hidden_states
- def prepare_attn_masks(
- self,
- input_ids: torch.Tensor,
- positions: torch.Tensor,
- mask_dtype: torch.dtype,
- **kwargs,
- ):
- kwargs["has_images"] = True
- # NOTE(woosuk): Here, we distinguish the sequences by the position id 0.
- # This is a HACK. Fix this.
- start_indices = (positions == 0).cpu().nonzero()
- num_seqs = len(start_indices)
- seq_lens = []
- for i in range(num_seqs):
- start_idx = start_indices[i].item()
- if i < num_seqs - 1:
- end_idx = start_indices[i + 1].item()
- else:
- end_idx = len(input_ids)
- seq_lens.append(end_idx - start_idx)
- kwargs["seq_lens"] = seq_lens
-
- global_attn_masks = []
- local_attn_masks = []
- start_idx = 0
- for seq_len in seq_lens:
- end_idx = start_idx + seq_len
- input_token_ids = input_ids[start_idx:end_idx]
- start_idx = end_idx
- # Create a global causal mask.
- global_attn_mask = torch.empty(
- 1,
- 1,
- seq_len,
- seq_len,
- dtype=mask_dtype,
- device=input_ids.device,
- )
- global_attn_mask.fill_(float("-inf"))
- # Fill the lower triangle with 0.
- global_attn_mask = global_attn_mask.triu(diagonal=1)
-
- # Consider the bidirectional attention between image tokens.
- img_mask = torch.zeros_like(global_attn_mask)
- img_pos = input_token_ids == self.config.image_token_index
- img_mask[:, :, :, img_pos] += 1
- img_mask[:, :, img_pos, :] += 1
- global_attn_mask = torch.where(img_mask == 2, 0, global_attn_mask)
- global_attn_masks.append(global_attn_mask)
-
- sliding_window = self.config.text_config.sliding_window
- if sliding_window is not None:
- # Create a local causal mask with sliding window (1024).
- local_attn_mask = torch.ones_like(global_attn_mask)
- local_attn_mask = torch.tril(local_attn_mask, diagonal=-sliding_window)
- local_attn_mask = torch.where(
- local_attn_mask == 0, global_attn_mask, float("-inf")
- )
- local_attn_masks.append(local_attn_mask)
- kwargs["global_attn_masks"] = global_attn_masks
- kwargs["local_attn_masks"] = local_attn_masks
- return kwargs
-
def compute_logits(
self,
hidden_states: torch.Tensor,
diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py
index 64443190f53ed..8f1447ba34a81 100644
--- a/vllm/model_executor/models/gemma3n.py
+++ b/vllm/model_executor/models/gemma3n.py
@@ -332,18 +332,21 @@ class Gemma3nAttention(nn.Module):
)
layer_idx = extract_layer_index(prefix)
- is_sliding = config.layer_types[layer_idx] == "sliding_attention"
+ layer_type = config.layer_types[layer_idx]
+ is_sliding = layer_type == "sliding_attention"
self.sliding_window = config.sliding_window if is_sliding else None
# Initialize the rotary embedding.
- if is_sliding:
- # Local attention. Override the values in config.json.
- rope_theta = config.rope_local_base_freq
- rope_scaling = {"rope_type": "default"}
+ if layer_type in config.rope_parameters:
+ # Transformers v5 rope config.
+ rope_parameters = config.rope_parameters[layer_type]
else:
+ # Transformers v4 rope config.
# Global attention. Use the values in config.json.
- rope_theta = config.rope_theta
- rope_scaling = config.rope_scaling
+ rope_parameters = config.rope_parameters.copy()
+ # Local attention. Override the values in config.json.
+ if is_sliding:
+ rope_parameters["rope_theta"] = config.rope_local_base_freq
first_kv_shared_layer_idx = (
config.num_hidden_layers - config.num_kv_shared_layers
@@ -383,9 +386,8 @@ class Gemma3nAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
+ rope_parameters=rope_parameters,
is_neox_style=True,
- rope_scaling=rope_scaling,
)
self.attn = Attention(
diff --git a/vllm/model_executor/models/glm4.py b/vllm/model_executor/models/glm4.py
index faa0674a2e43d..f8ef3b0385fb1 100644
--- a/vllm/model_executor/models/glm4.py
+++ b/vllm/model_executor/models/glm4.py
@@ -57,10 +57,8 @@ class Glm4Attention(nn.Module):
max_position: int = 4096 * 32,
head_dim: int | None = None,
qkv_bias: bool = False,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
- rope_scaling: tuple | None = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER,
) -> None:
@@ -86,7 +84,6 @@ class Glm4Attention(nn.Module):
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.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
@@ -107,8 +104,7 @@ class Glm4Attention(nn.Module):
self.head_dim,
rotary_dim=self.rotary_dim,
max_position=max_position,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
partial_rotary_factor=partial_rotary_factor,
is_neox_style=False,
)
@@ -150,8 +146,6 @@ class Glm4DecoderLayer(nn.Module):
quant_config = vllm_config.quant_config
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 1000000)
- rope_scaling = getattr(config, "rope_scaling", None)
self.self_attn = Glm4Attention(
config=config,
@@ -159,12 +153,10 @@ class Glm4DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
qkv_bias=getattr(config, "attention_bias", False),
head_dim=getattr(config, "head_dim", None),
cache_config=cache_config,
quant_config=quant_config,
- rope_scaling=rope_scaling,
prefix=f"{prefix}.self_attn",
attn_type=AttentionType.DECODER,
)
diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py
index 6953b805653b4..7e0370886884f 100644
--- a/vllm/model_executor/models/glm4_1v.py
+++ b/vllm/model_executor/models/glm4_1v.py
@@ -37,7 +37,7 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
-from transformers import BatchFeature
+from transformers import BatchFeature, Glm4vProcessor
from transformers.models.glm4v.configuration_glm4v import Glm4vVisionConfig
from transformers.models.glm4v.image_processing_glm4v import (
Glm4vImageProcessor,
@@ -56,7 +56,7 @@ from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
from vllm.distributed import utils as dist_utils
from vllm.logger import init_logger
-from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
@@ -65,6 +65,7 @@ from vllm.model_executor.layers.linear import (
RowParallelLinear,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -308,7 +309,6 @@ class Glm4vVisionAttention(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -341,9 +341,9 @@ class Glm4vVisionAttention(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: int | None = None, # Only used for Flash Attention
- seqlens: list[int] | None = None, # Only used for xFormers
) -> torch.Tensor:
# [s, b, c] --> [s, b, head * 3 * head_dim]
x, _ = self.qkv(x)
@@ -353,10 +353,12 @@ class Glm4vVisionAttention(nn.Module):
batch_size = q.shape[1]
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v))
- if rotary_pos_emb is not None:
+ if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
# [2 * b, s, heads, head_dim]
qk_concat = torch.cat([q, k], dim=0)
- qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
+ qk_rotated = apply_rotary_pos_emb_vision(
+ qk_concat, rotary_pos_emb_cos, rotary_pos_emb_sin
+ )
q, k = torch.chunk(qk_rotated, 2, dim=0)
if self.is_flash_attn_backend:
@@ -396,20 +398,6 @@ class Glm4vVisionAttention(nn.Module):
context_layer = rearrange(
context_layer, "b s h d -> s b (h d)"
).contiguous()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
-
- attn_bias = BlockDiagonalMask.from_seqlens(
- q_seqlen=seqlens, kv_seqlen=None, device=q.device
- )
-
- context_layer = xops.memory_efficient_attention_forward(
- q, k, v, attn_bias=attn_bias, p=0, scale=None
- )
- context_layer = rearrange(
- context_layer, "b s h d -> s b (h d)"
- ).contiguous()
output, _ = self.proj(context_layer)
return output
@@ -454,16 +442,16 @@ class Glm4vVisionBlock(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: int | None = None, # Only used for Flash Attention
- seqlens: list[int] | None = None, # Only used for xFormers
) -> torch.Tensor:
x_attn = self.attn(
self.norm1(x),
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
x_fused_norm, residual = self.norm2(x, residual=x_attn)
x = residual + self.mlp(x_fused_norm)
@@ -660,44 +648,6 @@ class Glm4vVisionEmbeddings(nn.Module):
return embeddings
-class Glm4vVisionRotaryEmbedding(nn.Module):
- def __init__(self, dim: int, theta: float = 10000.0) -> None:
- super().__init__()
- self.dim = dim
- self.theta = theta
- inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
- self._seq_len_cached = 0
- self._freqs_cached = None
-
- def update_freqs_cache(self, seqlen: int) -> None:
- if seqlen > self._seq_len_cached:
- seqlen *= 2
- self._seq_len_cached = seqlen
- self.inv_freq = 1.0 / (
- self.theta
- ** (
- torch.arange(
- 0,
- self.dim,
- 2,
- dtype=torch.float,
- device=self.inv_freq.device,
- )
- / self.dim
- )
- )
- seq = torch.arange(
- seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
- )
- freqs = torch.outer(seq, self.inv_freq)
- self._freqs_cached = freqs
-
- def forward(self, seqlen: int) -> torch.Tensor:
- self.update_freqs_cache(seqlen)
- return self._freqs_cached[:seqlen]
-
-
class Glm4vVisionTransformer(nn.Module):
def __init__(
self,
@@ -731,7 +681,12 @@ class Glm4vVisionTransformer(nn.Module):
norm_layer = partial(RMSNorm, eps=norm_eps)
head_dim = self.hidden_size // self.num_heads
- self.rotary_pos_emb = Glm4vVisionRotaryEmbedding(head_dim // 2)
+ self.rotary_pos_emb = get_rope(
+ head_size=head_dim,
+ rotary_dim=head_dim // 2,
+ max_position=8192,
+ is_neox_style=True,
+ )
self.blocks = nn.ModuleList(
[
Glm4vVisionBlock(
@@ -760,7 +715,7 @@ class Glm4vVisionTransformer(nn.Module):
self.post_conv_layernorm = RMSNorm(
vision_config.hidden_size, eps=vision_config.rms_norm_eps
)
- self.downsample = nn.Conv2d(
+ self.downsample = Conv2dLayer(
in_channels=vision_config.hidden_size,
out_channels=vision_config.out_hidden_size,
kernel_size=vision_config.spatial_merge_size,
@@ -789,7 +744,9 @@ class Glm4vVisionTransformer(nn.Module):
def device(self) -> torch.device:
return self.patch_embed.proj.weight.device
- def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+ def rot_pos_emb(
+ self, grid_thw: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
pos_ids = []
for t, h, w in grid_thw:
hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
@@ -817,22 +774,25 @@ class Glm4vVisionTransformer(nn.Module):
pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
pos_ids = torch.cat(pos_ids, dim=0)
max_grid_size = grid_thw[:, 1:].max()
- rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
- rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
- return rotary_pos_emb, pos_ids
+
+ # Use pre-computed cos_sin_cache from RotaryEmbedding
+ cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)
+
+ cos_combined = cos[pos_ids].flatten(1)
+ sin_combined = sin[pos_ids].flatten(1)
+ return cos_combined, sin_combined, pos_ids
def compute_attn_mask_seqlen(
self,
cu_seqlens: torch.Tensor,
- ) -> tuple[int | None, list[int] | None]:
- max_seqlen, seqlens = None, None
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+ ) -> int | None:
+ max_seqlen = None
if (
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
):
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
- return max_seqlen, seqlens
+ return max_seqlen
def forward(
self,
@@ -848,15 +808,18 @@ class Glm4vVisionTransformer(nn.Module):
x = self.post_conv_layernorm(x)
# compute position embedding
- rotary_pos_emb, image_type_ids = self.rot_pos_emb(grid_thw)
+ rotary_pos_emb_cos, rotary_pos_emb_sin, image_type_ids = self.rot_pos_emb(
+ grid_thw
+ )
# compute cu_seqlens
cu_seqlens = torch.repeat_interleave(
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
).cumsum(dim=0, dtype=torch.int32)
cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)
- # pre-compute seqlens for attn mask to reduce cuMemcpy operations
- max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+ # pre-compute max_seqlen for attn mask to reduce cuMemcpy operations
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
x = self.embeddings(
x, seqlens, grid_thw, image_type_ids[:, 0], image_type_ids[:, 1]
)
@@ -867,9 +830,9 @@ class Glm4vVisionTransformer(nn.Module):
x = blk(
x,
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
# adapter
@@ -1046,7 +1009,7 @@ class Glm4vProcessingInfo(BaseProcessingInfo):
return max(max_frames_per_video, 1)
- def _get_video_second_idx(
+ def _get_video_second_idx_glm4v(
self, metadata: dict[str, Any], total_frames: int
) -> list[int]:
video_processor = self.get_video_processor()
@@ -1097,6 +1060,83 @@ class Glm4vProcessingInfo(BaseProcessingInfo):
selected_timestamps.append(timestamps_list[idx])
return selected_timestamps
+ def _get_video_second_idx_glm46v(
+ self, metadata: dict[str, Any], total_frames: int
+ ) -> list[int]:
+ video_processor = self.get_video_processor()
+
+ video_fps = metadata["fps"]
+ meta_frames = metadata.get("total_num_frames", total_frames)
+ max_frame_idx = meta_frames - 1
+ duration = metadata.get("duration", round(max_frame_idx / video_fps) + 1)
+
+ do_sample_frames = metadata.get("do_sample_frames", True)
+ if not do_sample_frames:
+ frame_indices = metadata["frames_indices"]
+ else:
+ DYNAMIC_FPS_THRES = {30: 3, 300: 1, 2400: 0.5}
+ MAX_FRAME_COUNT_DYNAMIC = 640
+ MAX_DURATION = 2400
+
+ effective_duration = min(duration, MAX_DURATION)
+ if effective_duration <= 30:
+ target_fps = DYNAMIC_FPS_THRES[30]
+ elif effective_duration <= 300:
+ target_fps = DYNAMIC_FPS_THRES[300]
+ else:
+ target_fps = DYNAMIC_FPS_THRES[2400]
+
+ temporal_patch_size = getattr(video_processor, "temporal_patch_size", 1)
+ extract_t = int(effective_duration * target_fps * temporal_patch_size)
+ extract_t = min(extract_t, MAX_FRAME_COUNT_DYNAMIC)
+
+ duration_per_frame = 1 / video_fps
+ timestamps = [i * duration_per_frame for i in range(meta_frames)]
+ max_second = int(duration)
+
+ if meta_frames < extract_t:
+ frame_indices = np.linspace(
+ 0, meta_frames - 1, extract_t, dtype=int
+ ).tolist()
+ else:
+ frame_indices = []
+ current_second = 0.0
+ inv_fps = 1 / (temporal_patch_size * target_fps)
+ for frame_index in range(meta_frames):
+ if timestamps[frame_index] >= current_second:
+ current_second += inv_fps
+ frame_indices.append(frame_index)
+ if current_second >= max_second:
+ break
+
+ if len(frame_indices) < extract_t:
+ if len(frame_indices) == 0:
+ start, end = 0, max(meta_frames - 1, 0)
+ else:
+ start, end = frame_indices[0], frame_indices[-1]
+ frame_indices = np.linspace(start, end, extract_t, dtype=int).tolist()
+ elif len(frame_indices) > extract_t:
+ frame_indices = np.linspace(
+ 0, meta_frames - 1, extract_t, dtype=int
+ ).tolist()
+
+ seen, uniq = set(), []
+ for idx in frame_indices:
+ if idx not in seen:
+ seen.add(idx)
+ uniq.append(idx)
+
+ if len(uniq) & 1:
+ uniq.append(uniq[-1])
+
+ frame_indices = uniq
+ full_second_idxs = [int(idx / video_fps) for idx in frame_indices]
+ timestamps_list = full_second_idxs[::2]
+ selected_timestamps = []
+ for idx in range(len(timestamps_list)):
+ selected_timestamps.append(timestamps_list[idx])
+ return selected_timestamps
+
def _construct_video_placeholder(
self,
video_array: np.ndarray,
@@ -1115,9 +1155,18 @@ class Glm4vProcessingInfo(BaseProcessingInfo):
merge_length = image_processor.merge_size**2
assert isinstance(grid_thw, torch.Tensor)
- timestamps = self._get_video_second_idx(metadata, len(video_array))
+ timestamps = (
+ self._get_video_second_idx_glm4v(metadata, len(video_array))
+ if isinstance(hf_processor, Glm4vProcessor)
+ else self._get_video_second_idx_glm46v(metadata, len(video_array))
+ )
+
+ timestamp_format = (
+ "{}" if isinstance(hf_processor, Glm4vProcessor) else "{:.1f} seconds"
+ )
frames_idx_token = [
- tokenizer.encode(str(i), add_special_tokens=False) for i in timestamps
+ tokenizer.encode(timestamp_format.format(i), add_special_tokens=False)
+ for i in timestamps
]
T, H, W = grid_thw
num_tokens_per_frame = int(H * W) // merge_length
diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py
index 1422dbe9b3cd0..5aa51af54a00b 100644
--- a/vllm/model_executor/models/glm4_moe.py
+++ b/vllm/model_executor/models/glm4_moe.py
@@ -26,7 +26,6 @@
import typing
from collections.abc import Callable, Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -233,8 +232,6 @@ class Glm4MoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 131072,
head_dim: int | None = None,
rms_norm_eps: float = 1e-05,
@@ -264,7 +261,6 @@ class Glm4MoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.use_qk_norm = use_qk_norm
@@ -291,8 +287,7 @@ class Glm4MoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
partial_rotary_factor=partial_rotary_factor,
)
self.attn = Attention(
@@ -341,8 +336,6 @@ class Glm4MoeDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
# DecoderLayers are created with `make_layers` which passes the prefix
# with the layer's index.
@@ -354,8 +347,6 @@ class Glm4MoeDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
head_dim=config.head_dim,
rms_norm_eps=config.rms_norm_eps,
diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py
index 1c18ea0745f2b..514082cf60ce2 100644
--- a/vllm/model_executor/models/glm4v.py
+++ b/vllm/model_executor/models/glm4v.py
@@ -24,6 +24,7 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import SiluAndMul, get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
@@ -78,7 +79,7 @@ class GLMVImagePixelInputs(TensorSchema):
class EVA2CLIPPatchEmbedding(nn.Module):
def __init__(self, config):
super().__init__()
- self.proj = nn.Conv2d(
+ self.proj = Conv2dLayer(
config.in_channels,
config.hidden_size,
kernel_size=config.patch_size,
@@ -333,7 +334,7 @@ class EVA2CLIPModel(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.linear_proj",
)
- self.conv = nn.Conv2d(
+ self.conv = Conv2dLayer(
in_channels=vision_config.hidden_size,
out_channels=config.hidden_size,
kernel_size=2,
diff --git a/vllm/model_executor/models/gpt_j.py b/vllm/model_executor/models/gpt_j.py
index e416ecde0c1e0..e94de8952fa63 100644
--- a/vllm/model_executor/models/gpt_j.py
+++ b/vllm/model_executor/models/gpt_j.py
@@ -95,13 +95,12 @@ class GPTJAttention(nn.Module):
scaling = self.head_size**-0.5
assert getattr(config, "rotary", True)
assert config.rotary_dim % 2 == 0
- rope_theta = getattr(config, "rope_theta", 10000)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.rotary_emb = get_rope(
self.head_size,
rotary_dim=config.rotary_dim,
max_position=max_position_embeddings,
- base=rope_theta,
+ rope_parameters=config.rope_parameters,
is_neox_style=False,
)
self.attn = Attention(
diff --git a/vllm/model_executor/models/gpt_neox.py b/vllm/model_executor/models/gpt_neox.py
index af0c9209231cb..815c2fba4d9fe 100644
--- a/vllm/model_executor/models/gpt_neox.py
+++ b/vllm/model_executor/models/gpt_neox.py
@@ -92,13 +92,12 @@ class GPTNeoXAttention(nn.Module):
scaling = self.head_size**-0.5
rotary_dim = int(self.head_size * config.rotary_pct)
assert rotary_dim % 2 == 0
- rope_theta = getattr(config, "rope_theta", 10000)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.rotary_emb = get_rope(
self.head_size,
rotary_dim=rotary_dim,
max_position=max_position_embeddings,
- base=rope_theta,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py
index 7df3b087ccb88..1bc0ad38765d5 100644
--- a/vllm/model_executor/models/gpt_oss.py
+++ b/vllm/model_executor/models/gpt_oss.py
@@ -13,6 +13,7 @@ from vllm.config import CacheConfig, VllmConfig
from vllm.distributed import (
get_dp_group,
get_ep_group,
+ get_pcp_group,
get_pp_group,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
@@ -67,16 +68,17 @@ class OAIAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=config.max_position_embeddings,
- base=config.rope_theta,
dtype=torch.float32,
- rope_scaling={
+ rope_parameters={
+ "rope_theta": config.rope_parameters["rope_theta"],
"rope_type": "yarn",
- "factor": config.rope_scaling["factor"],
- "original_max_position_embeddings": config.rope_scaling[
+ "factor": config.rope_parameters["factor"],
+ "original_max_position_embeddings": config.rope_parameters[
"original_max_position_embeddings"
],
- "beta_fast": config.rope_scaling["beta_fast"],
- "beta_slow": config.rope_scaling["beta_slow"],
+ "beta_fast": config.rope_parameters["beta_fast"],
+ "beta_slow": config.rope_parameters["beta_slow"],
+ "truncate": config.rope_parameters.get("truncate", True),
},
is_neox_style=True,
)
@@ -90,7 +92,6 @@ class OAIAttention(nn.Module):
self.q_size = self.num_attention_heads * self.head_dim // tp_size
self.kv_size = self.num_key_value_heads * self.head_dim // tp_size
self.scaling = self.head_dim**-0.5
- self.rope_theta = config.rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size=self.hidden_size,
@@ -323,10 +324,12 @@ class GptOssModel(nn.Module):
# In MoE, we need to flatten the tensor parallel size across the data
# parallel size when EP is disabled.
- tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
+ tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
tp_size=get_tensor_model_parallel_world_size(),
dp_size=get_dp_group().world_size,
dp_rank=get_dp_group().rank_in_group,
+ pcp_size=get_pcp_group().world_size,
+ pcp_rank=get_pcp_group().rank_in_group,
)
intermediate_size = self.config.intermediate_size
@@ -508,10 +511,12 @@ class GptOssModel(nn.Module):
# In MoE, we need to flatten the tensor parallel size across the data
# parallel size when EP is disabled.
- tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
+ tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
tp_size=get_tensor_model_parallel_world_size(),
dp_size=get_dp_group().world_size,
dp_rank=get_dp_group().rank_in_group,
+ pcp_size=get_pcp_group().world_size,
+ pcp_rank=get_pcp_group().rank_in_group,
)
intermediate_size = self.config.intermediate_size
@@ -651,6 +656,7 @@ class GptOssModel(nn.Module):
class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
+ is_3d_moe_weight: bool = True
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
hf_to_vllm_mapper = WeightsMapper(
diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py
index c44b4021471ef..cd7ce2fc8f00a 100644
--- a/vllm/model_executor/models/granite.py
+++ b/vllm/model_executor/models/granite.py
@@ -26,7 +26,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -47,7 +46,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 (
- DEFAULT_VOCAB_PADDING_SIZE,
ParallelLMHead,
VocabParallelEmbedding,
)
@@ -112,8 +110,6 @@ class GraniteAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -143,7 +139,6 @@ class GraniteAttention(nn.Module):
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = config.attention_multiplier
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -167,8 +162,7 @@ class GraniteAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -204,14 +198,6 @@ class GraniteDecoderLayer(nn.Module):
super().__init__()
self.hidden_size = config.hidden_size
self.residual_multiplier = config.residual_multiplier
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -225,8 +211,6 @@ class GraniteDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
@@ -276,29 +260,16 @@ class GraniteModel(nn.Module):
config = vllm_config.model_config.hf_config
cache_config = vllm_config.cache_config
quant_config = vllm_config.quant_config
- lora_config = vllm_config.lora_config
self.config = config
self.quant_config = quant_config
- lora_vocab = (
- (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
- if lora_config
- else 0
- )
- self.vocab_size = config.vocab_size + lora_vocab
- self.org_vocab_size = config.vocab_size
+
if get_pp_group().is_first_rank or (
config.tie_word_embeddings and get_pp_group().is_last_rank
):
self.embed_tokens = VocabParallelEmbedding(
- self.vocab_size,
+ config.vocab_size,
config.hidden_size,
- org_num_embeddings=config.vocab_size,
- padding_size=DEFAULT_VOCAB_PADDING_SIZE
- # We need bigger padding if using lora for kernel
- # compatibility
- if not lora_config
- else lora_config.lora_vocab_padding_size,
quant_config=quant_config,
)
else:
@@ -435,28 +406,18 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
- lora_config = vllm_config.lora_config
self.config = config
- self.lora_config = lora_config
+
self.quant_config = quant_config
self.model = GraniteModel(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
if get_pp_group().is_last_rank:
- self.unpadded_vocab_size = config.vocab_size
- if lora_config:
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
self.lm_head = ParallelLMHead(
- self.unpadded_vocab_size,
+ config.vocab_size,
config.hidden_size,
- org_num_embeddings=config.vocab_size,
- padding_size=DEFAULT_VOCAB_PADDING_SIZE
- # We need bigger padding if using lora for kernel
- # compatibility
- if not lora_config
- else lora_config.lora_vocab_padding_size,
quant_config=quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
@@ -468,7 +429,7 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
logit_scale /= config.logits_scaling
self.logits_processor = LogitsProcessor(
- self.unpadded_vocab_size, config.vocab_size, scale=logit_scale
+ config.vocab_size, scale=logit_scale
)
else:
self.lm_head = PPMissingLayer()
diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py
index 5c6759ded0669..8f4139d63c3f6 100644
--- a/vllm/model_executor/models/granitemoe.py
+++ b/vllm/model_executor/models/granitemoe.py
@@ -141,8 +141,7 @@ class GraniteMoeAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
attention_multiplier: float | None = None,
@@ -172,7 +171,6 @@ class GraniteMoeAttention(nn.Module):
if attention_multiplier is not None
else self.head_dim**-1
)
- self.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -194,9 +192,8 @@ class GraniteMoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=int(self.rope_theta),
+ rope_parameters=rope_parameters,
is_neox_style=True,
- rope_scaling=rope_scaling,
)
self.attn = Attention(
self.num_heads,
@@ -235,16 +232,12 @@ class GraniteMoeDecoderLayer(nn.Module):
parallel_config = vllm_config.parallel_config
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
self.self_attn = GraniteMoeAttention(
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_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py
index 05177f1d1ac2c..9d5eeef198a61 100644
--- a/vllm/model_executor/models/granitemoehybrid.py
+++ b/vllm/model_executor/models/granitemoehybrid.py
@@ -115,8 +115,7 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
):
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
- output = torch.empty_like(hidden_states)
- self.mamba(hidden_states, output)
+ output = self.mamba(hidden_states)
hidden_states = residual + output * self.residual_multiplier
residual = hidden_states
@@ -274,10 +273,7 @@ class GraniteMoeHybridAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=config.max_position_embeddings,
- base=int(config.rope_theta),
- rope_scaling=config.rope_scaling
- if hasattr(config, "rope_scaling") and config.rope_scaling is not None
- else None,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
else:
diff --git a/vllm/model_executor/models/granitemoeshared.py b/vllm/model_executor/models/granitemoeshared.py
index 926c539af33be..fd346db7e35aa 100644
--- a/vllm/model_executor/models/granitemoeshared.py
+++ b/vllm/model_executor/models/granitemoeshared.py
@@ -84,16 +84,12 @@ class GraniteMoeSharedDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
self.self_attn = GraniteMoeAttention(
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_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
diff --git a/vllm/model_executor/models/grok1.py b/vllm/model_executor/models/grok1.py
index 9dc231863f74f..4bf23cd6fd19a 100644
--- a/vllm/model_executor/models/grok1.py
+++ b/vllm/model_executor/models/grok1.py
@@ -25,6 +25,7 @@
from collections.abc import Iterable
from itertools import islice
+from typing import Any
import torch
import torch.nn.functional as F
@@ -134,7 +135,7 @@ class Grok1Attention(nn.Module):
num_heads: int,
num_kv_heads: int,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
+ rope_parameters: dict[str, Any] | None = None,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
@@ -161,7 +162,6 @@ class Grok1Attention(nn.Module):
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.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -183,7 +183,7 @@ class Grok1Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=int(self.rope_theta),
+ rope_parameters=rope_parameters,
is_neox_style=True,
)
@@ -234,15 +234,12 @@ class Grok1DecoderLayer(nn.Module):
if not self.use_fp8 and hasattr(quant_config, "is_fp8"):
self.use_fp8 = quant_config.is_fp8
- # Requires transformers > 4.32.0
- # Default rope_theta value if not in config
- rope_theta = 10000
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_theta=rope_theta,
+ rope_parameters=config.rope_parameters,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.attn",
diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py
index 1eadcbe67ade3..9fa5e2bd33f21 100644
--- a/vllm/model_executor/models/hunyuan_v1.py
+++ b/vllm/model_executor/models/hunyuan_v1.py
@@ -27,7 +27,6 @@
import typing
from collections.abc import Callable, Iterable
from itertools import islice
-from typing import Any
import regex as re
import torch
@@ -142,8 +141,6 @@ class HunYuanAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -177,7 +174,6 @@ class HunYuanAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.use_qk_norm = getattr(config, "use_qk_norm", False)
self.layer_id = layer_id
@@ -204,8 +200,7 @@ class HunYuanAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
@@ -254,8 +249,6 @@ class HunYuanCrossAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -289,7 +282,6 @@ class HunYuanCrossAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.use_qk_norm = getattr(config, "use_qk_norm", False)
self.layer_id = layer_id
@@ -314,8 +306,7 @@ class HunYuanCrossAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
@@ -494,14 +485,6 @@ class HunYuanDecoderLayer(nn.Module):
if isinstance(config.intermediate_size, int)
else config.intermediate_size[layer_id]
)
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
attention_bias = getattr(config, "attention_bias", False) or getattr(
config, "bias", False
@@ -520,8 +503,6 @@ class HunYuanDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
@@ -537,8 +518,6 @@ class HunYuanDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py
index 727c8ec0397ca..06b8468e18db9 100644
--- a/vllm/model_executor/models/idefics2_vision_model.py
+++ b/vllm/model_executor/models/idefics2_vision_model.py
@@ -30,6 +30,7 @@ from transformers.models.idefics2.configuration_idefics2 import (
from vllm.attention.layer import MultiHeadAttention
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -60,7 +61,7 @@ class Idefics2VisionEmbeddings(nn.Module):
self.embed_dim = config.hidden_size
self.image_size = config.image_size
self.patch_size = config.patch_size
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=config.num_channels,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py
index dc4caf2f02f9d..9966498e1b4c9 100644
--- a/vllm/model_executor/models/interfaces.py
+++ b/vllm/model_executor/models/interfaces.py
@@ -586,13 +586,11 @@ class IsHybrid(Protocol):
def get_mamba_state_shape_from_config(
cls,
vllm_config: VllmConfig,
- use_v1: bool = True,
) -> tuple[tuple[int, int], tuple[int, int, int]]:
"""Calculate shapes for Mamba's convolutional and state caches.
Args:
vllm_config: vLLM config
- use_v1: Get shapes for V1 (or V0)
Returns:
Tuple containing:
diff --git a/vllm/model_executor/models/intern_vit.py b/vllm/model_executor/models/intern_vit.py
index 03918127c6ae1..61aeafc2ab436 100644
--- a/vllm/model_executor/models/intern_vit.py
+++ b/vllm/model_executor/models/intern_vit.py
@@ -24,6 +24,7 @@ from vllm.distributed import (
tensor_model_parallel_all_gather,
)
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
@@ -51,7 +52,7 @@ class InternVisionEmbeddings(nn.Module):
self.class_embedding = nn.Parameter(torch.randn(1, 1, self.embed_dim))
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=3,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
diff --git a/vllm/model_executor/models/internlm2.py b/vllm/model_executor/models/internlm2.py
index 60fbeb842dd4b..dc8f821bd134f 100644
--- a/vllm/model_executor/models/internlm2.py
+++ b/vllm/model_executor/models/internlm2.py
@@ -91,8 +91,7 @@ class InternLM2Attention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -120,7 +119,6 @@ class InternLM2Attention(nn.Module):
self.kv_size = self.num_kv_heads * self.head_dim
self.key_value_groups = int(self.num_heads / self.num_kv_heads)
self.scaling = self.head_dim**-0.5
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.wqkv = QKVParallelLinear(
@@ -144,8 +142,7 @@ class InternLM2Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -204,15 +201,12 @@ class InternLMDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.attention = InternLM2Attention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/internlm2_ve.py b/vllm/model_executor/models/internlm2_ve.py
index 6dc081e34157b..a57db82242af9 100644
--- a/vllm/model_executor/models/internlm2_ve.py
+++ b/vllm/model_executor/models/internlm2_ve.py
@@ -30,15 +30,12 @@ class InternLM2VEDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.attention = InternLM2Attention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/interns1_vit.py b/vllm/model_executor/models/interns1_vit.py
index 507503d75046d..cb0414bbc95a8 100644
--- a/vllm/model_executor/models/interns1_vit.py
+++ b/vllm/model_executor/models/interns1_vit.py
@@ -16,6 +16,7 @@ from transformers.utils import torch_int
from vllm.attention.layer import MultiHeadAttention
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -43,7 +44,7 @@ class InternS1VisionPatchEmbeddings(nn.Module):
self.num_patches = num_patches
self.patch_shape = patch_shape
- self.projection = nn.Conv2d(
+ self.projection = Conv2dLayer(
num_channels, hidden_size, kernel_size=patch_size, stride=patch_size
)
diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py
index 1eb0eccc0411c..302260b952992 100644
--- a/vllm/model_executor/models/keye.py
+++ b/vllm/model_executor/models/keye.py
@@ -9,6 +9,7 @@ from typing import Annotated, Any, Literal, TypeAlias, TypeVar
import numpy as np
import torch
import torch.nn as nn
+import torch.nn.functional as F
from einops import rearrange
from transformers import PretrainedConfig
from transformers.activations import GELUActivation
@@ -24,6 +25,7 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.logger import init_logger
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -204,7 +206,7 @@ class KeyeVisionEmbeddings(nn.Module):
self.image_size = config.image_size
self.patch_size = config.patch_size
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=config.num_channels,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
@@ -423,7 +425,7 @@ class KeyeSiglipAttention(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
- AttentionBackendEnum.XFORMERS,
+ AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -450,7 +452,6 @@ class KeyeSiglipAttention(nn.Module):
)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
batch_size = q.shape[0]
if rope_emb is None:
@@ -497,17 +498,21 @@ class KeyeSiglipAttention(nn.Module):
softmax_scale=self.scale,
)
context_layer = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
-
- attn_bias = BlockDiagonalMask.from_seqlens(
- q_seqlen=seqlens, kv_seqlen=None, device=q.device
- )
-
- context_layer = xops.memory_efficient_attention_forward(
- q, k, v, attn_bias=attn_bias, p=0, scale=None
- )
+ elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
+ outputs = []
+ for i in range(1, len(cu_seqlens)):
+ start_idx = cu_seqlens[i - 1]
+ end_idx = cu_seqlens[i]
+ q_i = q[:, start_idx:end_idx]
+ k_i = k[:, start_idx:end_idx]
+ v_i = v[:, start_idx:end_idx]
+ q_i, k_i, v_i = (
+ rearrange(x, "b s h d -> b h s d") for x in (q_i, k_i, v_i)
+ )
+ output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
+ output_i = rearrange(output_i, "b h s d -> b s h d ")
+ outputs.append(output_i)
+ context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
context_layer = rearrange(context_layer, "b s h d -> b s (h d)").contiguous()
diff --git a/vllm/model_executor/models/kimi_linear.py b/vllm/model_executor/models/kimi_linear.py
index f3675075a48f4..4562b2202c5ec 100644
--- a/vllm/model_executor/models/kimi_linear.py
+++ b/vllm/model_executor/models/kimi_linear.py
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
-from typing import Any
import torch
from torch import nn
@@ -190,9 +189,7 @@ class KimiMLAAttention(nn.Module):
v_head_dim: int,
q_lora_rank: int | None,
kv_lora_rank: int,
- rope_theta: float = 10000,
use_nope: bool = False,
- rope_scaling: dict[str, Any] | None = None,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
@@ -210,11 +207,9 @@ class KimiMLAAttention(nn.Module):
tp_size = get_tensor_model_parallel_world_size()
self.num_local_heads = num_heads // tp_size
self.scaling = self.qk_head_dim**-0.5
- self.rope_theta = rope_theta
self.use_nope = use_nope
assert self.use_nope is True
assert self.q_lora_rank is None
- assert rope_scaling is None
assert num_heads % tp_size == 0
self.kv_a_proj_with_mqa = ReplicatedLinear(
self.hidden_size,
diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py
index aeb25602f11a4..74bdde27ece5c 100644
--- a/vllm/model_executor/models/lfm2.py
+++ b/vllm/model_executor/models/lfm2.py
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
import torch.nn as nn
@@ -96,8 +95,6 @@ class Lfm2Attention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -126,7 +123,6 @@ class Lfm2Attention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -149,8 +145,7 @@ class Lfm2Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
@@ -199,14 +194,6 @@ class Lfm2AttentionDecoderLayer(nn.Module):
self.config = config
self.layer_idx = layer_idx
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.self_attn = Lfm2Attention(
@@ -215,8 +202,6 @@ class Lfm2AttentionDecoderLayer(nn.Module):
hidden_size=config.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py
index 6b7b5564ee989..c088a08211527 100644
--- a/vllm/model_executor/models/lfm2_moe.py
+++ b/vllm/model_executor/models/lfm2_moe.py
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
import torch.nn as nn
@@ -189,8 +188,6 @@ class Lfm2MoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -219,7 +216,6 @@ class Lfm2MoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -242,8 +238,7 @@ class Lfm2MoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
@@ -293,14 +288,6 @@ class Lfm2MoeAttentionDecoderLayer(nn.Module):
self.config = config
self.layer_idx = layer_idx
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.self_attn = Lfm2MoeAttention(
@@ -309,8 +296,6 @@ class Lfm2MoeAttentionDecoderLayer(nn.Module):
hidden_size=config.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
index 0a3f37c30ab5f..ebf8addda4a54 100644
--- a/vllm/model_executor/models/llama.py
+++ b/vllm/model_executor/models/llama.py
@@ -26,7 +26,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -48,7 +47,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 (
- DEFAULT_VOCAB_PADDING_SIZE,
ParallelLMHead,
VocabParallelEmbedding,
)
@@ -120,8 +118,6 @@ class LlamaAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -157,7 +153,6 @@ class LlamaAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
@@ -186,9 +181,7 @@ class LlamaAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
- self._init_rotary_emb(
- config, rope_scaling=rope_scaling, quant_config=quant_config
- )
+ self._init_rotary_emb(config, quant_config=quant_config)
sliding_window = None
if layer_types := getattr(config, "layer_types", None):
@@ -258,7 +251,6 @@ class LlamaAttention(nn.Module):
def _init_rotary_emb(
self,
config: LlamaConfig,
- rope_scaling: dict[str, Any] | None,
quant_config: QuantizationConfig | None,
) -> None:
is_neox_style = True
@@ -270,8 +262,7 @@ class LlamaAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
partial_rotary_factor=self.partial_rotary_factor,
)
@@ -291,14 +282,6 @@ class LlamaDecoderLayer(nn.Module):
quant_config = self.get_quant_config(vllm_config)
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -326,8 +309,6 @@ class LlamaDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
@@ -386,24 +367,18 @@ class LlamaModel(nn.Module):
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
- lora_config = vllm_config.lora_config
self.config = config
self.quant_config = quant_config
- lora_vocab = (
- (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
- if lora_config
- else 0
- )
- self.vocab_size = config.vocab_size + lora_vocab
- self.org_vocab_size = config.vocab_size
+
+ self.vocab_size = config.vocab_size
+
if get_pp_group().is_first_rank or (
config.tie_word_embeddings and get_pp_group().is_last_rank
):
self.embed_tokens = VocabParallelEmbedding(
self.vocab_size,
config.hidden_size,
- org_num_embeddings=config.vocab_size,
quant_config=quant_config,
)
else:
@@ -580,9 +555,7 @@ class LlamaForCausalLM(
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
- lora_config = vllm_config.lora_config
self.config = config
- self.lora_config = lora_config
self.model = self._init_model(
vllm_config=vllm_config,
@@ -591,20 +564,9 @@ class LlamaForCausalLM(
)
if get_pp_group().is_last_rank:
- self.unpadded_vocab_size = config.vocab_size
- if lora_config:
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
self.lm_head = ParallelLMHead(
- self.unpadded_vocab_size,
+ config.vocab_size,
config.hidden_size,
- org_num_embeddings=config.vocab_size,
- padding_size=(
- DEFAULT_VOCAB_PADDING_SIZE
- # We need bigger padding if using lora for kernel
- # compatibility
- if not lora_config
- else lora_config.lora_vocab_padding_size
- ),
quant_config=quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
@@ -613,7 +575,7 @@ class LlamaForCausalLM(
logit_scale = getattr(config, "logit_scale", 1.0)
self.logits_processor = LogitsProcessor(
- self.unpadded_vocab_size, config.vocab_size, logit_scale
+ config.vocab_size, scale=logit_scale
)
else:
self.lm_head = PPMissingLayer()
diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py
index a7e0732ec71e2..e1bdfc3405f70 100644
--- a/vllm/model_executor/models/llama4.py
+++ b/vllm/model_executor/models/llama4.py
@@ -19,7 +19,6 @@
"""Inference-only LLaMA model compatible with HuggingFace weights."""
from collections.abc import Iterable
-from typing import Any
import torch
from torch import nn
@@ -54,6 +53,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
from .utils import (
AutoWeightsLoader,
+ PPMissingLayer,
extract_layer_index,
fast_topk,
is_pp_missing_parameter,
@@ -171,8 +171,6 @@ class Llama4Attention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -208,7 +206,6 @@ class Llama4Attention(nn.Module):
self.floor_scale = getattr(config, "floor_scale", 8192.0)
self.attn_scale = getattr(config, "attn_scale", 0.1)
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.n_rep = self.num_heads // self.num_kv_heads
self.qk_norm = (
@@ -248,8 +245,7 @@ class Llama4Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=int(rope_theta),
- rope_scaling=rope_scaling if rope_scaling != "default" else None,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
)
if not self.nope
@@ -331,8 +327,6 @@ class Llama4DecoderLayer(nn.Module):
self.layer_idx = extract_layer_index(prefix)
self.global_layer = config.no_rope_layers[self.layer_idx] == 0
self.hidden_size = config.hidden_size
- rope_theta = config.rope_theta
- rope_scaling = config.rope_scaling
max_position_embeddings = config.max_position_embeddings
self.self_attn = Llama4Attention(
@@ -340,8 +334,6 @@ class Llama4DecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=False,
@@ -738,6 +730,9 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
self.moe_layers = []
example_moe = None
for layer in self.model.layers:
+ if isinstance(layer, PPMissingLayer):
+ continue
+
assert isinstance(layer, Llama4DecoderLayer)
if isinstance(layer.feed_forward, Llama4MoE):
# Pick last one layer since the first ones may be dense layers.
@@ -774,6 +769,9 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
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 isinstance(layer.feed_forward, Llama4MoE):
moe = layer.feed_forward
moe.n_local_physical_experts = num_local_physical_experts
diff --git a/vllm/model_executor/models/llama4_eagle.py b/vllm/model_executor/models/llama4_eagle.py
index 660c8f1bb5226..0146b30579287 100644
--- a/vllm/model_executor/models/llama4_eagle.py
+++ b/vllm/model_executor/models/llama4_eagle.py
@@ -23,7 +23,6 @@ import torch.nn as nn
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
-from vllm.distributed.parallel_state import get_pp_group
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -127,17 +126,11 @@ class LlamaModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
- # if PP disabled then draft will share embed with target
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
- continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
for name in params_dict:
- # if PP disabled then draft will share embed with target
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
- continue
assert name in loaded_params, f"{name} is not loaded!"
return loaded_params
diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py
index 90ab5c50361b6..05cb456e7776e 100644
--- a/vllm/model_executor/models/llama_eagle.py
+++ b/vllm/model_executor/models/llama_eagle.py
@@ -9,7 +9,6 @@ from transformers import LlamaConfig
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
-from vllm.distributed.parallel_state import get_pp_group
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -155,10 +154,6 @@ class LlamaModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
- # if PP disabled then draft will share embed with target
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
- continue
-
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py
index 75c671311b491..3eaf2d80082f1 100644
--- a/vllm/model_executor/models/llama_eagle3.py
+++ b/vllm/model_executor/models/llama_eagle3.py
@@ -23,7 +23,6 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM
-from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import NestedTensors
from .utils import (
@@ -121,13 +120,12 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
@support_torch_compile(
- # torch.compile is disabled for multimodal EAGLE3 models due to constraint
- # violations with dynamic shapes during tensor concatenation operations.
- # See: https://github.com/vllm-project/vllm/pull/22872/files#r2362028132
- # Non-multimodal EAGLE3 models can still use torch.compile safely.
- enable_if=lambda vllm_config: not MULTIMODAL_REGISTRY.supports_multimodal_inputs(
- vllm_config.model_config
- ),
+ dynamic_arg_dims={
+ "input_ids": 0,
+ "positions": -1,
+ "hidden_states": 0,
+ "input_embeds": 0,
+ }
)
class LlamaModel(nn.Module):
def __init__(
diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py
index 5de10e7086830..c5441283f9711 100644
--- a/vllm/model_executor/models/longcat_flash.py
+++ b/vllm/model_executor/models/longcat_flash.py
@@ -108,8 +108,7 @@ class FlashConfig(PretrainedConfig):
eos_token_id=100001,
pretraining_tp=1,
tie_word_embeddings=False,
- rope_theta=1000000.0,
- rope_scaling=None,
+ rope_parameters=None,
attention_bias=False,
attention_dropout=0.0,
mla_scale_q_lora=False,
@@ -119,7 +118,7 @@ class FlashConfig(PretrainedConfig):
router_dtype="float32",
router_bias=False,
topk_method=None,
- routed_scaling_factor=None,
+ routed_scaling_factor=1.0,
zero_expert_num=0,
zero_expert_type=None,
nextn_use_scmoe=False,
@@ -162,8 +161,13 @@ class FlashConfig(PretrainedConfig):
self.rms_norm_eps = rms_norm_eps
self.pretraining_tp = pretraining_tp
self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 1000000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
self.mla_scale_q_lora = mla_scale_q_lora
@@ -336,15 +340,7 @@ class FlashDecoderLayer(nn.Module):
super().__init__()
self.layer_idx = int(prefix.split(sep=".")[-1])
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
# Dual attention structure
self.self_attn = nn.ModuleList(
@@ -361,8 +357,6 @@ class FlashDecoderLayer(nn.Module):
config.q_lora_rank if hasattr(config, "q_lora_rank") else None
),
kv_lora_rank=config.kv_lora_rank,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=None
diff --git a/vllm/model_executor/models/mamba2.py b/vllm/model_executor/models/mamba2.py
index fc17f98be1986..5fcfa94312303 100644
--- a/vllm/model_executor/models/mamba2.py
+++ b/vllm/model_executor/models/mamba2.py
@@ -87,8 +87,7 @@ class Mamba2DecoderLayer(nn.Module):
else:
hidden_states, residual = self.norm(hidden_states, residual)
- output = torch.empty_like(hidden_states)
- self.mixer(hidden_states, output)
+ output = self.mixer(hidden_states)
return output, residual
diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py
index a84c99059cd9c..d9b23811730d4 100644
--- a/vllm/model_executor/models/midashenglm.py
+++ b/vllm/model_executor/models/midashenglm.py
@@ -39,6 +39,7 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -120,7 +121,7 @@ class AudioPatchEmbed(nn.Module):
self.num_patches = self.grid_size[0] * self.grid_size[1]
self.flatten = flatten
- self.proj = nn.Conv2d(
+ self.proj = Conv2dLayer(
in_chans,
embed_dim,
kernel_size=self.patch_size,
diff --git a/vllm/model_executor/models/minicpm.py b/vllm/model_executor/models/minicpm.py
index 914b097fe199e..04923833065f3 100644
--- a/vllm/model_executor/models/minicpm.py
+++ b/vllm/model_executor/models/minicpm.py
@@ -230,8 +230,7 @@ class MiniCPMAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -257,7 +256,6 @@ class MiniCPMAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -281,8 +279,7 @@ class MiniCPMAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
@@ -324,8 +321,6 @@ class MiniCPMDecoderLayer(nn.Module):
self.cache_config = cache_config
self.quant_config = quant_config
self.hidden_size = config.hidden_size
- self.rope_theta = getattr(config, "rope_theta", 10000)
- self.rope_scaling = getattr(config, "rope_scaling", None)
self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.prefix = prefix
self._init_attn_block()
@@ -339,8 +334,7 @@ class MiniCPMDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=self.config.num_attention_heads,
num_kv_heads=self.config.num_key_value_heads,
- rope_theta=self.rope_theta,
- rope_scaling=self.rope_scaling,
+ rope_parameters=self.config.rope_parameters,
max_position_embeddings=self.max_position_embeddings,
cache_config=self.cache_config,
quant_config=self.quant_config,
diff --git a/vllm/model_executor/models/minicpm3.py b/vllm/model_executor/models/minicpm3.py
index d3b6966ee3a7f..2d775219fc972 100644
--- a/vllm/model_executor/models/minicpm3.py
+++ b/vllm/model_executor/models/minicpm3.py
@@ -25,8 +25,6 @@
# limitations under the License.
"""Inference-only MiniCPM3 model compatible with HuggingFace weights."""
-from typing import Any
-
import torch
from torch import nn
from transformers import PretrainedConfig
@@ -62,8 +60,6 @@ class MiniCPM3Attention(nn.Module):
v_head_dim: int,
q_lora_rank: int,
kv_lora_rank: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -84,7 +80,6 @@ class MiniCPM3Attention(nn.Module):
self.num_local_heads = num_heads // tp_size
self.scaling = self.qk_head_dim**-0.5
- self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.q_a_proj = ReplicatedLinear(
@@ -127,8 +122,7 @@ class MiniCPM3Attention(nn.Module):
self.qk_rope_head_dim,
rotary_dim=self.qk_rope_head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_local_heads,
@@ -204,8 +198,6 @@ class MiniCPM3DecoderLayer(MiniCPMDecoderLayer):
v_head_dim=self.config.v_head_dim,
q_lora_rank=self.config.q_lora_rank,
kv_lora_rank=self.config.kv_lora_rank,
- rope_theta=self.rope_theta,
- rope_scaling=self.rope_scaling,
max_position_embeddings=self.max_position_embeddings,
cache_config=self.cache_config,
quant_config=self.quant_config,
diff --git a/vllm/model_executor/models/minicpm_eagle.py b/vllm/model_executor/models/minicpm_eagle.py
index d0cdb70aa8574..e6bccfcac4f1a 100644
--- a/vllm/model_executor/models/minicpm_eagle.py
+++ b/vllm/model_executor/models/minicpm_eagle.py
@@ -69,8 +69,6 @@ class EagleMiniCPMDecoderLayer(nn.Module):
self.cache_config = cache_config
self.quant_config = quant_config
self.hidden_size = config.hidden_size
- self.rope_theta = getattr(config, "rope_theta", 10000)
- self.rope_scaling = getattr(config, "rope_scaling", None)
self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.prefix = prefix
self._init_attn_block()
@@ -84,8 +82,7 @@ class EagleMiniCPMDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=self.config.num_attention_heads,
num_kv_heads=self.config.num_key_value_heads,
- rope_theta=self.rope_theta,
- rope_scaling=self.rope_scaling,
+ rope_parameters=self.config.rope_parameters,
max_position_embeddings=self.max_position_embeddings,
cache_config=self.cache_config,
quant_config=self.quant_config,
diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py
index 49d2f2d261969..4955c68c0cda8 100644
--- a/vllm/model_executor/models/minimax_m2.py
+++ b/vllm/model_executor/models/minimax_m2.py
@@ -149,8 +149,7 @@ class MiniMaxM2Attention(nn.Module):
num_heads: int,
num_kv_heads: int,
rotary_dim: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
attn_window_size: int | None = None,
max_position_embeddings: int = 8192,
head_dim: int | None = None,
@@ -180,7 +179,6 @@ class MiniMaxM2Attention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -205,8 +203,7 @@ class MiniMaxM2Attention(nn.Module):
self.head_dim,
rotary_dim=rotary_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -252,8 +249,6 @@ class MiniMaxM2DecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
if hasattr(config, "max_model_len") and isinstance(config.max_model_len, int):
max_position_embeddings = max(
@@ -269,8 +264,7 @@ class MiniMaxM2DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
rotary_dim=config.rotary_dim,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
rms_norm_eps=config.rms_norm_eps,
qkv_bias=getattr(config, "attention_bias", False),
diff --git a/vllm/model_executor/models/minimax_text_01.py b/vllm/model_executor/models/minimax_text_01.py
index bf1ecc822756d..50f7396e2de60 100644
--- a/vllm/model_executor/models/minimax_text_01.py
+++ b/vllm/model_executor/models/minimax_text_01.py
@@ -188,7 +188,7 @@ class MiniMaxText01Attention(nn.Module):
num_kv_heads: int,
rotary_dim: int,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
+ rope_parameters: dict | None = None,
sliding_window: int | None = None,
quant_config: QuantizationConfig | None = None,
layer_idx: int = None,
@@ -214,7 +214,6 @@ class MiniMaxText01Attention(nn.Module):
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.rope_theta = rope_theta
self.sliding_window = sliding_window
self.prefix = prefix
@@ -247,7 +246,7 @@ class MiniMaxText01Attention(nn.Module):
head_size=self.head_dim,
rotary_dim=rotary_dim,
max_position=max_position,
- base=int(rope_theta),
+ rope_parameters=rope_parameters,
is_neox_style=True,
dtype=torch.float32,
)
@@ -287,8 +286,6 @@ class MiniMaxText01DecoderLayer(nn.Module):
self.hidden_size = config.hidden_size
self.expert_num = expert_num
- rope_theta = getattr(config, "rope_theta", 10000)
-
head_dim = getattr(config, "head_dim", None)
if head_dim is None:
head_dim = config.hidden_size // config.num_attention_heads
@@ -328,7 +325,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
else head_dim,
num_kv_heads=config.num_key_value_heads,
max_position=max_position_embeddings,
- rope_theta=rope_theta,
+ rope_parameters=config.rope_parameters,
sliding_window=config.sliding_window,
quant_config=quant_config,
layer_idx=self._ilayer,
diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py
index d7a1cb82fb4fb..0a9c3f136964e 100644
--- a/vllm/model_executor/models/mixtral.py
+++ b/vllm/model_executor/models/mixtral.py
@@ -51,7 +51,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 (
- DEFAULT_VOCAB_PADDING_SIZE,
ParallelLMHead,
VocabParallelEmbedding,
)
@@ -161,7 +160,6 @@ class MixtralAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
@@ -189,7 +187,6 @@ class MixtralAttention(nn.Module):
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.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -211,7 +208,7 @@ class MixtralAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=int(self.rope_theta),
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
@@ -248,15 +245,12 @@ class MixtralDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 10000)
self.self_attn = MixtralAttention(
config=config,
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_theta=rope_theta,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
@@ -306,23 +300,18 @@ class MixtralModel(nn.Module):
config = vllm_config.model_config.hf_config
cache_config = vllm_config.cache_config
quant_config = vllm_config.quant_config
- lora_config = vllm_config.lora_config
+
parallel_config = vllm_config.parallel_config
self.config = config
self.quant_config = quant_config
- lora_vocab = (
- (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
- if lora_config
- else 0
- )
- self.vocab_size = config.vocab_size + lora_vocab
+
+ self.vocab_size = config.vocab_size
self.org_vocab_size = config.vocab_size
self.embed_tokens = VocabParallelEmbedding(
self.vocab_size,
config.hidden_size,
- org_num_embeddings=config.vocab_size,
)
self.enable_eplb = parallel_config.enable_eplb
@@ -513,34 +502,24 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
- lora_config = vllm_config.lora_config
+
self.config = config
- self.lora_config = lora_config
+
self.quant_config = quant_config
self.model = MixtralModel(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
- self.unpadded_vocab_size = config.vocab_size
- if lora_config:
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
+
self.lm_head = ParallelLMHead(
- self.unpadded_vocab_size,
+ config.vocab_size,
config.hidden_size,
- org_num_embeddings=config.vocab_size,
- padding_size=DEFAULT_VOCAB_PADDING_SIZE
- # We need bigger padding if using lora for kernel
- # compatibility
- if not lora_config
- else lora_config.lora_vocab_padding_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.logits_processor = LogitsProcessor(
- self.unpadded_vocab_size, config.vocab_size
- )
+ self.logits_processor = LogitsProcessor(config.vocab_size)
self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors
)
diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py
index e25a104d822a7..286859d188d34 100644
--- a/vllm/model_executor/models/mllama4.py
+++ b/vllm/model_executor/models/mllama4.py
@@ -292,13 +292,17 @@ class Llama4VisionAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
+ rope_parameters = {
+ "rope_type": "mllama4",
+ "rope_theta": config.rope_parameters["rope_theta"],
+ }
+
self.rotary_emb = get_rope(
head_size=self.head_dim,
rotary_dim=config.hidden_size // config.num_attention_heads // 2,
# number of image patches
max_position=(config.image_size // config.patch_size) ** 2,
- base=config.rope_theta,
- rope_scaling={"rope_type": "mllama4"},
+ rope_parameters=rope_parameters,
is_neox_style=False,
dtype=torch.complex64, # important
)
diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py
index ab83a271e30a0..dc06938d5d6e1 100644
--- a/vllm/model_executor/models/molmo.py
+++ b/vllm/model_executor/models/molmo.py
@@ -410,7 +410,6 @@ class MolmoAttention(nn.Module):
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
# Attention input projection. Projects x -> (q, k, v)
self.qkv_proj = QKVParallelLinear(
@@ -437,7 +436,7 @@ class MolmoAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=config.rope_parameters,
)
self.scaling = self.head_dim**-0.5
self.attn = Attention(
diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py
index 8017c947bf9ad..2e3e6dc166ad8 100644
--- a/vllm/model_executor/models/moonvit.py
+++ b/vllm/model_executor/models/moonvit.py
@@ -53,6 +53,7 @@ from transformers.activations import ACT2FN
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import is_flash_attn_2_available
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.models.utils import maybe_prefix
from vllm.transformers_utils.configs.moonvit import MoonViTConfig
@@ -244,7 +245,7 @@ class MoonVisionPatchEmbed(nn.Module):
)
self.patch_size = patch_size
- self.proj = nn.Conv2d(
+ self.proj = Conv2dLayer(
in_dim, out_dim, kernel_size=patch_size, stride=patch_size
)
diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py
index 92dcf5ea57008..c3337bd1ea699 100644
--- a/vllm/model_executor/models/nemotron.py
+++ b/vllm/model_executor/models/nemotron.py
@@ -26,7 +26,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -150,8 +149,6 @@ class NemotronAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -181,7 +178,6 @@ class NemotronAttention(nn.Module):
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.rope_theta = rope_theta
self.partial_rotary_factor = config.partial_rotary_factor
self.max_position_embeddings = max_position_embeddings
@@ -206,8 +202,7 @@ class NemotronAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
partial_rotary_factor=self.partial_rotary_factor,
)
self.attn = Attention(
@@ -243,14 +238,6 @@ class NemotronDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -264,8 +251,6 @@ class NemotronDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py
index f7e0caf410e10..8675eff592224 100644
--- a/vllm/model_executor/models/nemotron_h.py
+++ b/vllm/model_executor/models/nemotron_h.py
@@ -376,8 +376,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
else:
hidden_states, residual = self.norm(hidden_states, residual)
- output = torch.empty_like(hidden_states)
- self.mixer(hidden_states, output)
+ output = self.mixer(hidden_states)
return output, residual
diff --git a/vllm/model_executor/models/nemotron_nas.py b/vllm/model_executor/models/nemotron_nas.py
index b839206a3094d..2eebe38051cbd 100644
--- a/vllm/model_executor/models/nemotron_nas.py
+++ b/vllm/model_executor/models/nemotron_nas.py
@@ -26,7 +26,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import Any
import torch
from torch import nn
@@ -82,8 +81,6 @@ class DeciLMAttention(LlamaAttention):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -97,8 +94,6 @@ class DeciLMAttention(LlamaAttention):
hidden_size,
num_heads,
num_kv_heads,
- rope_theta,
- rope_scaling,
max_position_embeddings,
quant_config,
bias,
@@ -111,7 +106,6 @@ class DeciLMAttention(LlamaAttention):
def _init_rotary_emb(
self,
config,
- rope_scaling: dict[str, Any] | None,
quant_config: QuantizationConfig | None,
) -> None:
# Enables YARN for Mistral and LLaMA4 derivatives.
@@ -126,8 +120,7 @@ class DeciLMAttention(LlamaAttention):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
partial_rotary_factor=self.partial_rotary_factor,
)
@@ -148,14 +141,6 @@ class DeciLMDecoderLayer(nn.Module):
self._is_no_op_ffn = block_config.ffn.no_op
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -176,8 +161,6 @@ class DeciLMDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=num_kv_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/olmo.py b/vllm/model_executor/models/olmo.py
index 487e3f671a455..bd8a8e317544f 100644
--- a/vllm/model_executor/models/olmo.py
+++ b/vllm/model_executor/models/olmo.py
@@ -87,7 +87,6 @@ class OlmoAttention(nn.Module):
self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
self.head_dim = self.hidden_size // self.total_num_heads
self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
self.clip_qkv = config.clip_qkv
# Attention input projection. Projects x -> (q, k, v)
@@ -105,7 +104,7 @@ class OlmoAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=config.rope_parameters,
)
self.scaling = self.head_dim**-0.5
self.attn = Attention(
diff --git a/vllm/model_executor/models/olmo2.py b/vllm/model_executor/models/olmo2.py
index 045582c889ee4..f0f6b2f6b3e6d 100644
--- a/vllm/model_executor/models/olmo2.py
+++ b/vllm/model_executor/models/olmo2.py
@@ -99,7 +99,6 @@ class Olmo2Attention(nn.Module):
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.max_position_embeddings = self.config.max_position_embeddings
- self.rope_theta = self.config.rope_theta
# Attention input projection. Projects x -> (q, k, v)
self.qkv_proj = QKVParallelLinear(
@@ -139,15 +138,17 @@ class Olmo2Attention(nn.Module):
prefix=f"{prefix}.attn",
)
- # Rotary embeddings. Rope scaling is only applied on full attention
- # layers.
- self.rope_scaling = self.config.rope_scaling if sliding_window is None else None
+ # Rotary embeddings. Rope scaling is only applied on full attention layers.
+ if sliding_window is None:
+ rope_parameters = self.config.rope_parameters
+ else:
+ rope_theta = self.config.rope_parameters["rope_theta"]
+ rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta, # type: ignore
- rope_scaling=self.rope_scaling,
+ rope_parameters=rope_parameters,
)
# Attention output projection.
diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py
index 499eb05de76e4..c39e338d72e22 100644
--- a/vllm/model_executor/models/olmoe.py
+++ b/vllm/model_executor/models/olmoe.py
@@ -123,8 +123,6 @@ class OlmoeAttention(nn.Module):
quant_config = vllm_config.quant_config
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 4096)
num_heads = config.num_attention_heads
@@ -148,7 +146,6 @@ class OlmoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -176,8 +173,7 @@ class OlmoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
diff --git a/vllm/model_executor/models/opencua.py b/vllm/model_executor/models/opencua.py
new file mode 100644
index 0000000000000..121bf896fa6ba
--- /dev/null
+++ b/vllm/model_executor/models/opencua.py
@@ -0,0 +1,271 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from Qwen2.5-VL implementation
+# Copyright 2025 The vLLM team.
+# Copyright 2025 XLANG Lab, The University of Hong Kong
+
+"""Inference-only OpenCUA-7B model compatible with HuggingFace weights."""
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import torch
+import torch.nn as nn
+from transformers import BatchFeature
+from transformers.models.qwen2_vl import (
+ Qwen2VLImageProcessor,
+ Qwen2VLProcessor,
+ Qwen2VLVideoProcessor,
+)
+
+from vllm.config import VllmConfig
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+ MultiModalFieldConfig,
+ MultiModalKwargs,
+)
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+ BaseMultiModalProcessor,
+ PromptReplacement,
+ PromptUpdate,
+)
+from vllm.transformers_utils.tokenizer import AnyTokenizer
+
+from .qwen2_5_vl import (
+ Qwen2_5_VisionTransformer as OpenCUAVisionTransformer,
+)
+from .qwen2_5_vl import (
+ Qwen2_5_VLForConditionalGeneration,
+)
+from .qwen2_vl import (
+ Qwen2VLDummyInputsBuilder,
+ Qwen2VLMultiModalDataParser,
+ Qwen2VLProcessingInfo,
+ _create_qwen2vl_field_factory,
+)
+from .utils import (
+ WeightsMapper,
+ init_vllm_registered_model,
+ maybe_prefix,
+)
+
+
+class OpenCUAProcessingInfo(Qwen2VLProcessingInfo):
+ def get_hf_config(self):
+ return self.ctx.get_hf_config()
+
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+ return {"image": None}
+
+ def get_hf_processor(self, **kwargs: object):
+ """Load OpenCUA processor."""
+ tokenizer = self.get_tokenizer()
+ vision_config = self.ctx.get_hf_image_processor_config()
+ return OpenCUAProcessor(
+ vision_config=vision_config,
+ tokenizer=tokenizer,
+ **kwargs,
+ )
+
+
+class OpenCUAProcessor(Qwen2VLProcessor):
+ def check_argument_for_proper_class(self, attribute_name: str, arg: object) -> None:
+ if attribute_name == "tokenizer":
+ return
+ return super().check_argument_for_proper_class(attribute_name, arg)
+
+ def __init__(
+ self,
+ vision_config: dict,
+ tokenizer: AnyTokenizer,
+ **kwargs,
+ ):
+ image_processor = Qwen2VLImageProcessor(**vision_config)
+ video_processor = Qwen2VLVideoProcessor(**vision_config)
+ chat_template = kwargs.pop("chat_template", None)
+
+ super().__init__(
+ image_processor=image_processor,
+ tokenizer=tokenizer,
+ video_processor=video_processor,
+ chat_template=chat_template,
+ **kwargs,
+ )
+
+ self.image_token = "<|media_placeholder|>"
+
+ def __call__(
+ self,
+ text=None,
+ images=None,
+ return_tensors=None,
+ **kwargs,
+ ):
+ if text is not None:
+ if not isinstance(text, list):
+ text = [text]
+ text_inputs = self.tokenizer(text, **kwargs)
+ else:
+ text_inputs = {}
+
+ image_inputs = {}
+ if images is not None:
+ if not isinstance(images, list):
+ images = [images]
+ if len(images) > 0:
+ image_inputs = self.image_processor(
+ images, return_tensors=return_tensors or "pt"
+ )
+
+ combined_inputs = {**text_inputs, **image_inputs}
+
+ return BatchFeature(combined_inputs, tensor_type=return_tensors)
+
+
+class OpenCUAMultiModalProcessor(BaseMultiModalProcessor[OpenCUAProcessingInfo]):
+ def _get_data_parser(self) -> MultiModalDataParser:
+ return Qwen2VLMultiModalDataParser(
+ self.info.get_hf_config().vision_config.spatial_merge_size
+ )
+
+ def _get_mm_fields_config(
+ self,
+ hf_inputs: BatchFeature,
+ hf_processor_mm_kwargs: Mapping[str, object],
+ ) -> Mapping[str, MultiModalFieldConfig]:
+ return _create_qwen2vl_field_factory(
+ self.info.get_hf_config().vision_config.spatial_merge_size
+ )(hf_inputs)
+
+ def _hf_processor_applies_updates(
+ self,
+ prompt_text: str,
+ mm_items: MultiModalDataItems,
+ hf_processor_mm_kwargs: Mapping[str, object],
+ tokenization_kwargs: Mapping[str, object],
+ ) -> bool:
+ """vLLM이 prompt 업데이트를 처리하도록 False 반환."""
+ return False
+
+ def _get_prompt_updates(
+ self,
+ mm_items: MultiModalDataItems,
+ hf_processor_mm_kwargs: Mapping[str, Any],
+ out_mm_kwargs: MultiModalKwargs,
+ ) -> Sequence[PromptUpdate]:
+ hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+ image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
+ tokenizer = self.info.get_tokenizer()
+ vocab = tokenizer.get_vocab()
+ hf_config = self.info.get_hf_config()
+
+ image_token_str = getattr(hf_processor, "image_token", "<|media_placeholder|>")
+ image_token_id = vocab.get(
+ image_token_str,
+ getattr(hf_config, "media_placeholder_token_id", 151664),
+ )
+
+ merge_length = image_processor.merge_size**2
+
+ def get_replacement_opencua(item_idx: int):
+ out_item = out_mm_kwargs["image"][item_idx]
+ grid_thw = out_item["image_grid_thw"].data
+ assert isinstance(grid_thw, torch.Tensor)
+
+ num_tokens = int(grid_thw.prod()) // merge_length
+ return [image_token_id] * num_tokens
+
+ return [
+ PromptReplacement(
+ modality="image",
+ target=[image_token_id],
+ replacement=get_replacement_opencua,
+ )
+ ]
+
+
+class OpenCUADummyInputsBuilder(Qwen2VLDummyInputsBuilder):
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+ num_images = mm_counts.get("image", 0)
+
+ image_token = "<|media_placeholder|>"
+
+ return image_token * num_images
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+ OpenCUAMultiModalProcessor,
+ info=OpenCUAProcessingInfo,
+ dummy_inputs=OpenCUADummyInputsBuilder,
+)
+class OpenCUAForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
+ merge_by_field_config = True
+ multimodal_cpu_fields = {"image_grid_thw"}
+
+ packed_modules_mapping = {
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+ "gate_up_proj": ["gate_proj", "up_proj"],
+ }
+
+ hf_to_vllm_mapper = WeightsMapper(
+ orig_to_new_prefix={
+ "model.language_model.": "language_model.model.",
+ "model.visual.": "visual.",
+ "vision_tower.": "visual.",
+ "lm_head.": "language_model.lm_head.",
+ "model.": "language_model.model.",
+ }
+ )
+
+ supports_encoder_tp_data = True
+
+ @classmethod
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+ if modality.startswith("image"):
+ return "<|media_placeholder|>"
+ raise ValueError("Only image modality is supported")
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ nn.Module.__init__(self)
+ config = vllm_config.model_config.hf_config
+ quant_config = vllm_config.quant_config
+ multimodal_config = vllm_config.model_config.multimodal_config
+
+ self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+ self.config = config
+ self.vllm_config = vllm_config
+ self.multimodal_config = multimodal_config
+ self.quant_config = quant_config
+ self.is_multimodal_pruning_enabled = (
+ multimodal_config.is_multimodal_pruning_enabled()
+ )
+
+ if multimodal_config.get_limit_per_prompt("image"):
+ attn_backend_override = (
+ multimodal_config.mm_encoder_attn_backend
+ if multimodal_config is not None
+ else None
+ )
+ self.visual = OpenCUAVisionTransformer(
+ vision_config=config.vision_config,
+ norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+ quant_config=self.quant_config,
+ prefix=maybe_prefix(prefix, "visual"),
+ use_data_parallel=self.use_data_parallel,
+ attn_backend_override=attn_backend_override,
+ )
+ else:
+ self.visual = None
+
+ self.language_model = init_vllm_registered_model(
+ vllm_config=vllm_config,
+ hf_config=config.text_config,
+ prefix=maybe_prefix(prefix, "language_model"),
+ architectures=["Qwen2ForCausalLM"],
+ )
+
+ self.make_empty_intermediate_tensors = (
+ self.language_model.make_empty_intermediate_tensors
+ )
diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py
index d13a745beffeb..4124a181a14c2 100644
--- a/vllm/model_executor/models/openpangu.py
+++ b/vllm/model_executor/models/openpangu.py
@@ -77,6 +77,7 @@ from vllm.model_executor.models.utils import (
sequence_parallel_chunk,
)
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
def check_ffn_act_fn(act_fn: str):
@@ -259,7 +260,6 @@ class OpenPanguMLAAttention(nn.Module):
v_head_dim: int,
q_lora_rank: int | None,
kv_lora_rank: int,
- rope_theta: float = 10000,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -274,8 +274,6 @@ class OpenPanguMLAAttention(nn.Module):
self.v_head_dim = v_head_dim
self.q_lora_rank = q_lora_rank
self.kv_lora_rank = kv_lora_rank
- self.rope_theta = rope_theta
-
self.tp_size = get_tensor_model_parallel_world_size()
if num_heads % self.tp_size != 0:
raise ValueError(
@@ -339,7 +337,9 @@ class OpenPanguMLAAttention(nn.Module):
)
# TODO: remove hard coding
- rope_scaling = {
+ set_default_rope_theta(config, default_theta=10000)
+ rope_parameters = {
+ "rope_theta": config.rope_parameters["rope_theta"],
"beta_fast": 32,
"beta_slow": 1,
"factor": 1,
@@ -353,8 +353,7 @@ class OpenPanguMLAAttention(nn.Module):
qk_rope_head_dim,
rotary_dim=qk_rope_head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
is_neox_style=False,
)
@@ -407,8 +406,6 @@ class OpenPanguEmbeddedAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -454,7 +451,6 @@ class OpenPanguEmbeddedAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -475,9 +471,7 @@ class OpenPanguEmbeddedAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
- self._init_rotary_emb(
- config, rope_scaling=rope_scaling, quant_config=quant_config
- )
+ self._init_rotary_emb(config, quant_config=quant_config)
if hasattr(config, "interleaved_sliding_window"):
interleaved_sliding_window = config.interleaved_sliding_window
@@ -521,7 +515,6 @@ class OpenPanguEmbeddedAttention(nn.Module):
def _init_rotary_emb(
self,
config: PretrainedConfig,
- rope_scaling: dict[str, Any] | None,
quant_config: QuantizationConfig | None,
) -> None:
is_neox_style = True
@@ -533,8 +526,7 @@ class OpenPanguEmbeddedAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
is_neox_style=is_neox_style,
)
@@ -555,7 +547,6 @@ class OpenPanguDecoderLayer(nn.Module):
parallel_config = vllm_config.parallel_config
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
layer_idx = int(prefix.split(sep=".")[-1])
@@ -579,7 +570,6 @@ class OpenPanguDecoderLayer(nn.Module):
config.q_lora_rank if hasattr(config, "q_lora_rank") else None
),
kv_lora_rank=config.kv_lora_rank,
- rope_theta=rope_theta,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
@@ -607,8 +597,6 @@ class OpenPanguDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=getattr(config, "rope_scaling", None),
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
@@ -637,7 +625,7 @@ class OpenPanguDecoderLayer(nn.Module):
bias=getattr(config, "mlp_bias", False),
prefix=f"{prefix}.mlp",
)
- self.routed_scaling_factor = getattr(config, "routed_scaling_factor", None)
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
self.num_hidden_layers = config.num_hidden_layers
self.first_k_dense_replace = getattr(
config, "first_k_dense_replace", self.num_hidden_layers
diff --git a/vllm/model_executor/models/orion.py b/vllm/model_executor/models/orion.py
index 859cd2cecf897..b30be93ca726f 100644
--- a/vllm/model_executor/models/orion.py
+++ b/vllm/model_executor/models/orion.py
@@ -88,8 +88,7 @@ class OrionAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -115,7 +114,6 @@ class OrionAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -139,8 +137,7 @@ class OrionAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -175,15 +172,12 @@ class OrionDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
self.self_attn = OrionAttention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/ouro.py b/vllm/model_executor/models/ouro.py
index 9db6c317c26a8..63d2fff6ec8bc 100644
--- a/vllm/model_executor/models/ouro.py
+++ b/vllm/model_executor/models/ouro.py
@@ -112,10 +112,8 @@ class OuroAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
- rope_scaling: tuple | None = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER,
dual_chunk_attention_config: dict[str, Any] | None = None,
@@ -140,7 +138,6 @@ class OuroAttention(nn.Module):
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.rope_theta = rope_theta
self.dual_chunk_attention_config = dual_chunk_attention_config
# Get total_ut_steps from config, default to 4 if not specified
@@ -170,8 +167,7 @@ class OuroAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
dual_chunk_attention_config=dual_chunk_attention_config,
)
self.attn = nn.ModuleList()
@@ -226,9 +222,6 @@ class OuroDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 1000000)
- rope_scaling = getattr(config, "rope_scaling", None)
dual_chunk_attention_config = getattr(
config, "dual_chunk_attention_config", None
)
@@ -244,10 +237,8 @@ class OuroDecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
cache_config=cache_config,
quant_config=quant_config,
- rope_scaling=rope_scaling,
prefix=f"{prefix}.self_attn",
attn_type=attn_type,
dual_chunk_attention_config=dual_chunk_attention_config,
diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py
index 3ef6470070d18..74bb868492da9 100644
--- a/vllm/model_executor/models/paddleocr_vl.py
+++ b/vllm/model_executor/models/paddleocr_vl.py
@@ -38,13 +38,13 @@ from vllm.attention.layer import (
)
from vllm.attention.ops.vit_attn_wrappers import (
vit_flash_attn_wrapper,
- vit_xformers_attn_wrapper,
)
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import parallel_state
from vllm.distributed import utils as dist_utils
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -419,7 +419,7 @@ class SiglipVisionEmbeddings(nn.Module):
self.image_size = config.image_size
self.patch_size = config.patch_size
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=config.num_channels,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
@@ -656,7 +656,6 @@ class SiglipAttention(nn.Module):
cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor | None,
max_seqlen: torch.Tensor | None,
- seqlens: torch.Tensor | None,
) -> torch.Tensor:
batch_size, _, _ = hidden_states.shape
@@ -702,10 +701,6 @@ class SiglipAttention(nn.Module):
context_layer = rearrange(
context_layer, "b s h d -> s b (h d)"
).contiguous()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- if seqlens is None:
- raise ValueError("xFormers attention backend requires seqlens tensor.")
- context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
else:
raise RuntimeError(
f"PaddleOCR-VL does not support {self.attn_backend} backend now."
@@ -817,7 +812,6 @@ class SiglipEncoderLayer(nn.Module):
cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor | None,
max_seqlen: torch.Tensor | None,
- seqlens: torch.Tensor | None,
) -> torch.Tensor:
residual = hidden_states
@@ -827,7 +821,6 @@ class SiglipEncoderLayer(nn.Module):
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
hidden_states = residual + hidden_states
@@ -869,7 +862,6 @@ class SiglipEncoder(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -942,14 +934,11 @@ class SiglipEncoder(nn.Module):
cu_seqlens = cu_seqlens.to(device=device)
max_seqlen = None
- seqlens = None
if self.attn_backend in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.ROCM_AITER_FA,
}:
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
hidden_states = inputs_embeds
for encoder_layer in self.layers:
@@ -958,7 +947,6 @@ class SiglipEncoder(nn.Module):
cu_seqlens=cu_seqlens,
rotary_pos_emb=rotary_pos_emb,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
return hidden_states
diff --git a/vllm/model_executor/models/persimmon.py b/vllm/model_executor/models/persimmon.py
index 3bf6a1d9763d0..98963d52e4848 100644
--- a/vllm/model_executor/models/persimmon.py
+++ b/vllm/model_executor/models/persimmon.py
@@ -106,7 +106,6 @@ class PersimmonAttention(nn.Module):
self.num_heads = self.total_num_heads // tensor_parallel_world_size
self.head_dim = self.hidden_size // self.total_num_heads
self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
self.partial_rotary_factor = config.partial_rotary_factor
self.is_causal = True
@@ -138,7 +137,7 @@ class PersimmonAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=self.rope_theta,
+ rope_parameters=config.rope_parameters,
partial_rotary_factor=self.partial_rotary_factor,
)
self.scaling = self.head_dim**-0.5
diff --git a/vllm/model_executor/models/phi.py b/vllm/model_executor/models/phi.py
index 8fee53c23fb4b..da476f621627b 100644
--- a/vllm/model_executor/models/phi.py
+++ b/vllm/model_executor/models/phi.py
@@ -115,16 +115,12 @@ class PhiAttention(nn.Module):
)
assert rotary_dim % 2 == 0
- # pylint: disable=C0301
- # Refer to:
- # https://huggingface.co/microsoft/phi-1_5/blob/d212a789620c380ff32ca1d1ee9943a777360987/modeling_phi.py#L518
- rope_theta = getattr(config, "rope_theta", 10000.0)
max_position_embeddings = getattr(config, "max_position_embeddings", 2048)
self.rotary_emb = get_rope(
self.head_size,
rotary_dim=rotary_dim,
max_position=max_position_embeddings,
- base=rope_theta,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
diff --git a/vllm/model_executor/models/phimoe.py b/vllm/model_executor/models/phimoe.py
index 92fd858b608bc..8ffac95d93960 100644
--- a/vllm/model_executor/models/phimoe.py
+++ b/vllm/model_executor/models/phimoe.py
@@ -86,7 +86,7 @@ class PhiMoEConfig(PretrainedConfig):
bos_token_id=1,
eos_token_id=2,
tie_word_embeddings=False,
- rope_theta=1e6,
+ rope_parameters=None,
sliding_window=None,
attention_dropout=0.0,
num_experts_per_tok=2,
@@ -119,7 +119,9 @@ class PhiMoEConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
+ if rope_parameters is None:
+ rope_theta = kwargs.pop("rope_theta", 1e6)
+ rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
self.attention_dropout = attention_dropout
self.num_experts_per_tok = num_experts_per_tok
@@ -302,12 +304,11 @@ class PhiMoEAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
+ rope_parameters: dict,
head_dim: int | None = None,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
- rope_scaling: dict | None = None,
prefix: str = "",
) -> None:
super().__init__()
@@ -332,8 +333,6 @@ class PhiMoEAttention(nn.Module):
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.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -355,9 +354,8 @@ class PhiMoEAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=int(self.rope_theta),
+ rope_parameters=rope_parameters,
is_neox_style=True,
- rope_scaling=self.rope_scaling,
)
self.attn = Attention(
self.num_heads,
@@ -393,7 +391,6 @@ class PhiMoEDecoderLayer(nn.Module):
super().__init__()
self.hidden_size = config.hidden_size
# Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 10000)
self.self_attn = PhiMoEAttention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
@@ -402,10 +399,9 @@ class PhiMoEDecoderLayer(nn.Module):
head_dim=getattr(
config, "head_dim", self.hidden_size // config.num_attention_heads
),
- rope_theta=rope_theta,
cache_config=cache_config,
quant_config=quant_config,
- rope_scaling=config.rope_scaling,
+ rope_parameters=config.rope_parameters,
prefix=f"{prefix}.self_attn",
)
self.block_sparse_moe = PhiMoE(
diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py
index 8cb7d6a889da4..6011d93a795d1 100644
--- a/vllm/model_executor/models/pixtral.py
+++ b/vllm/model_executor/models/pixtral.py
@@ -31,6 +31,7 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_and_mul_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
@@ -73,6 +74,7 @@ from .vision import (
)
try:
+ # Note: vLLM does not install xformers by default.
from xformers import ops as xops
if current_platform.is_cuda() and current_platform.has_device_capability(100):
@@ -747,7 +749,7 @@ class VisionTransformer(nn.Module):
def __init__(self, args: VisionEncoderArgs):
super().__init__()
self.args = args
- self.patch_conv = nn.Conv2d(
+ self.patch_conv = Conv2dLayer(
in_channels=args.num_channels,
out_channels=args.hidden_size,
kernel_size=args.patch_size,
@@ -1212,7 +1214,7 @@ class PixtralHFVisionModel(nn.Module):
self.config = config
- self.patch_conv = nn.Conv2d(
+ self.patch_conv = Conv2dLayer(
in_channels=config.num_channels,
out_channels=config.hidden_size,
kernel_size=config.patch_size,
diff --git a/vllm/model_executor/models/plamo2.py b/vllm/model_executor/models/plamo2.py
index 0c87f5000ff45..472de5590dcf8 100644
--- a/vllm/model_executor/models/plamo2.py
+++ b/vllm/model_executor/models/plamo2.py
@@ -4,10 +4,6 @@
from collections.abc import Iterable
from itertools import islice
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
- from vllm.attention.backends.abstract import AttentionBackend
import torch
from torch import nn
@@ -294,7 +290,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
has_decode = num_decodes > 0
num_actual_tokens = num_prefill_tokens + num_decodes
- # NOTE: V0 put prefill before decode, v1 puts decode before prefill
# Separate prefill and decode by splitting varlen input
# Split along token dimension
hidden_states_d, hidden_states_p = torch.split(
@@ -467,11 +462,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
def mamba_type(self) -> str:
return "mamba2"
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
-
- return Mamba2AttentionBackend
-
def plamo2_mamba_mixer(
hidden_states: torch.Tensor,
@@ -576,10 +566,6 @@ class Plamo2AttentionMixer(nn.Module):
prefix=f"{prefix}.o_proj",
)
- self.rope_theta = config.rope_theta if hasattr(config, "rope_theta") else 10000
- self.rope_scaling = (
- config.rope_scaling if hasattr(config, "rope_scaling") else None
- )
max_position = config.max_position_embeddings
if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
vllm_config.model_config.max_model_len, int
@@ -590,8 +576,7 @@ class Plamo2AttentionMixer(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=self.rope_theta,
- rope_scaling=self.rope_scaling,
+ rope_parameters=config.rope_parameters,
)
self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps)
self.q_norm.weight = torch.nn.Parameter(
diff --git a/vllm/model_executor/models/plamo3.py b/vllm/model_executor/models/plamo3.py
new file mode 100644
index 0000000000000..4aeb9d432dcc6
--- /dev/null
+++ b/vllm/model_executor/models/plamo3.py
@@ -0,0 +1,441 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only PLaMo3 model."""
+
+from collections.abc import Iterable
+from itertools import islice
+from typing import Any
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from vllm.attention.layer import Attention
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+ MergedColumnParallelLinear,
+ QKVParallelLinear,
+ 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 (
+ DEFAULT_VOCAB_PADDING_SIZE,
+ ParallelLMHead,
+ VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+ LoaderFunction,
+ composed_weight_loader,
+ default_weight_loader,
+)
+from vllm.model_executor.models.interfaces import SupportsPP
+from vllm.model_executor.models.utils import (
+ AutoWeightsLoader,
+ extract_layer_index,
+ make_empty_intermediate_tensors_factory,
+ make_layers,
+ maybe_prefix,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.sequence import IntermediateTensors
+
+
+# Only used for type hinting.
+class Plamo3Config(PretrainedConfig): # type: ignore
+ model_type: str = "plamo3"
+
+ hidden_size: int
+ num_hidden_layers: int
+ rms_norm_eps: float
+ # Attention
+ num_attention_heads: int
+ head_dim: int
+ num_key_value_heads: int
+ # vllm rename `sliding_window` attr to `interleaved_sliding_window`
+ # if `sliding_window` is list
+ interleaved_sliding_window: list[int | None]
+ sliding_window_pattern: int
+ rope_parameters: dict[str, Any]
+ rope_local_theta: int
+ # MLP
+ intermediate_size: int
+ # Tokenizer
+ vocab_size: int
+
+
+def rms_norm_weight_loader(offset: float) -> LoaderFunction:
+ return composed_weight_loader(
+ default_weight_loader,
+ lambda x: x + offset,
+ )
+
+
+class DenseMLP(nn.Module):
+ def __init__(
+ self,
+ config: Plamo3Config,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_up_proj = MergedColumnParallelLinear(
+ self.hidden_size,
+ [self.intermediate_size] * 2,
+ bias=False,
+ prefix=f"{prefix}.gate_up_proj",
+ quant_config=quant_config,
+ return_bias=False,
+ )
+ self.act = SiluAndMul()
+ self.down_proj = RowParallelLinear(
+ self.intermediate_size,
+ self.hidden_size,
+ bias=False,
+ prefix=f"{prefix}.down_proj",
+ quant_config=quant_config,
+ return_bias=False,
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ h = self.gate_up_proj(hidden_states)
+ h = self.act(h)
+ return self.down_proj(h)
+
+
+class Plamo3AttentionMixer(nn.Module):
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
+ super().__init__()
+ config = vllm_config.model_config.hf_config
+ quant_config = vllm_config.quant_config
+
+ self.hidden_size = config.hidden_size
+ tp_size = get_tensor_model_parallel_world_size()
+ self.total_num_heads = config.num_attention_heads
+ assert self.total_num_heads % tp_size == 0
+ self.num_heads = self.total_num_heads // tp_size
+ self.total_num_kv_heads = config.num_key_value_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 = config.head_dim
+ 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(
+ config.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,
+ config.hidden_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.o_proj",
+ )
+
+ layer_idx = extract_layer_index(prefix)
+ layer_type = config.layer_types[layer_idx]
+ is_sliding = layer_type == "sliding_attention"
+
+ # Initialize the rotary embedding.
+ if layer_type in config.rope_parameters:
+ # Transformers v5 rope config.
+ rope_parameters = config.rope_parameters[layer_type]
+ else:
+ # Transformers v4 rope config.
+ # Global attention. Use the values in config.json.
+ rope_parameters = config.rope_parameters
+ # Local attention. Override the values in config.json.
+ if is_sliding:
+ rope_parameters = dict(
+ rope_type="default", rope_theta=config.rope_local_theta
+ )
+ max_position = config.max_position_embeddings
+ if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
+ vllm_config.model_config.max_model_len, int
+ ):
+ max_position = min(max_position, vllm_config.model_config.max_model_len)
+
+ self.rotary_emb = get_rope(
+ self.head_dim,
+ rotary_dim=self.head_dim,
+ max_position=max_position,
+ rope_parameters=rope_parameters,
+ )
+ self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.q_norm.weight, {"weight_loader": rms_norm_weight_loader(offset=1.0)}
+ )
+ self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.k_norm.weight, {"weight_loader": rms_norm_weight_loader(offset=1.0)}
+ )
+ self.attn = Attention(
+ self.num_heads,
+ self.head_dim,
+ self.scaling,
+ num_kv_heads=self.num_kv_heads,
+ cache_config=vllm_config.cache_config,
+ per_layer_sliding_window=config.interleaved_sliding_window[layer_idx],
+ prefix=f"{prefix}.attn",
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ residual: torch.Tensor | None,
+ **kwargs: Any,
+ ) -> 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_shape = q.shape
+ q = q.reshape(q_shape[:-1] + (q_shape[-1] // self.head_dim, self.head_dim))
+ q = self.q_norm.forward_native(q).reshape(q_shape)
+ k_shape = k.shape
+ k = k.reshape(k_shape[:-1] + (k_shape[-1] // self.head_dim, self.head_dim))
+ k = self.k_norm.forward_native(k).reshape(k_shape)
+
+ q, k = self.rotary_emb(positions, q, k)
+ attn_output = self.attn(q, k, v)
+ output, _ = self.o_proj(attn_output)
+ return output
+
+
+class Plamo3DecoderLayer(nn.Module):
+ def __init__(
+ self, vllm_config: VllmConfig, prefix: str = "", **kwargs: Any
+ ) -> None:
+ super().__init__()
+ config = vllm_config.model_config.hf_config
+ quant_config = vllm_config.quant_config
+
+ self.mixer = Plamo3AttentionMixer(
+ vllm_config=vllm_config,
+ prefix=f"{prefix}.mixer",
+ )
+
+ self.mlp = DenseMLP(
+ config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
+ )
+ self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.pre_mixer_norm.weight,
+ {"weight_loader": rms_norm_weight_loader(offset=1.0)},
+ )
+ self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.post_mixer_norm.weight,
+ {"weight_loader": rms_norm_weight_loader(offset=1.0 / 5)},
+ )
+ self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.pre_mlp_norm.weight,
+ {"weight_loader": rms_norm_weight_loader(offset=1.0)},
+ )
+ self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.post_mlp_norm.weight,
+ {"weight_loader": rms_norm_weight_loader(offset=1.0 / (5**1.5))},
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ residual: torch.Tensor | None,
+ **kwargs: Any,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ if residual is None:
+ residual = hidden_states
+ hidden_states = self.pre_mixer_norm(hidden_states)
+ else:
+ hidden_states, residual = self.pre_mixer_norm(hidden_states, residual)
+
+ hidden_states = self.mixer(
+ positions=positions, hidden_states=hidden_states, residual=residual
+ )
+ hidden_states = self.post_mixer_norm(hidden_states)
+ # Fully Connected
+ hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.post_mlp_norm(hidden_states)
+ return hidden_states, residual
+
+
+class Plamo3Decoder(torch.nn.Module):
+ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+ num_hidden_layers = vllm_config.model_config.hf_config.num_hidden_layers
+
+ self.start_layer, self.end_layer, self.layers = make_layers(
+ num_hidden_layers,
+ lambda prefix: Plamo3DecoderLayer(vllm_config, prefix=prefix),
+ prefix=f"{prefix}.layers",
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ residual: torch.Tensor | None,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
+ hidden_states, residual = layer(
+ positions=positions,
+ hidden_states=hidden_states,
+ residual=residual,
+ )
+ return hidden_states, residual
+
+
+@support_torch_compile
+class Plamo3Model(nn.Module):
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+ config = vllm_config.model_config.hf_config
+
+ self.config = config
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self.org_vocab_size = config.vocab_size
+
+ self.embed_tokens = VocabParallelEmbedding(
+ self.vocab_size,
+ config.hidden_size,
+ org_num_embeddings=config.vocab_size,
+ prefix=f"{prefix}.embed_tokens",
+ )
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+ ["hidden_states", "residual"], config.hidden_size
+ )
+ self.layers = Plamo3Decoder(vllm_config, prefix=f"{prefix}.layers")
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ set_weight_attrs(
+ self.norm.weight,
+ {"weight_loader": rms_norm_weight_loader(offset=1.0)},
+ )
+
+ 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,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ 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"]
+
+ hidden_states, residual = self.layers(
+ positions=positions, hidden_states=hidden_states, residual=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
+
+
+class Plamo3ForCausalLM(nn.Module, SupportsPP):
+ packed_modules_mapping = {
+ "qkv_proj": [
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ ],
+ }
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+ self.config = vllm_config.model_config.hf_config
+ self.vllm_config = vllm_config
+ self.model_config = vllm_config.model_config
+ self.scheduler_config = vllm_config.scheduler_config
+
+ self.model = Plamo3Model(
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+ )
+
+ self.vocab_size = self.config.vocab_size
+ self.unpadded_vocab_size = self.config.vocab_size
+
+ num_embeddings = ((self.vocab_size + 15) // 16) * 16
+ self.lm_head = ParallelLMHead(
+ num_embeddings,
+ self.config.hidden_size,
+ org_num_embeddings=self.config.vocab_size,
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE,
+ prefix=f"{prefix}.lm_head",
+ )
+ if self.config.tie_word_embeddings:
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
+
+ self.logits_processor = LogitsProcessor(
+ self.unpadded_vocab_size, self.config.vocab_size
+ )
+ 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,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ 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]]):
+ loader = AutoWeightsLoader(
+ self,
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
+ )
+ return loader.load_weights(weights)
diff --git a/vllm/model_executor/models/qwen.py b/vllm/model_executor/models/qwen.py
index 50a125c3f5973..c973e79170982 100644
--- a/vllm/model_executor/models/qwen.py
+++ b/vllm/model_executor/models/qwen.py
@@ -83,8 +83,7 @@ class QWenAttention(nn.Module):
hidden_size: int,
num_heads: int,
max_position_embeddings: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
@@ -117,8 +116,7 @@ class QWenAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -153,14 +151,11 @@ class QWenBlock(nn.Module):
super().__init__()
self.ln_1 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
self.attn = QWenAttention(
config.hidden_size,
config.num_attention_heads,
config.max_position_embeddings,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.attn",
diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py
index 1bbb969ce5aa3..32b6d6dd07b83 100644
--- a/vllm/model_executor/models/qwen2.py
+++ b/vllm/model_executor/models/qwen2.py
@@ -57,7 +57,7 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.sequence import IntermediateTensors
-from vllm.transformers_utils.config import is_interleaved
+from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
from .utils import (
@@ -114,11 +114,10 @@ class Qwen2Attention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
+ rope_parameters: dict[str, Any],
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
- rope_scaling: tuple | None = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER,
dual_chunk_attention_config: dict[str, Any] | None = None,
@@ -143,7 +142,6 @@ class Qwen2Attention(nn.Module):
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.rope_theta = rope_theta
self.dual_chunk_attention_config = dual_chunk_attention_config
self.qkv_proj = QKVParallelLinear(
@@ -167,8 +165,7 @@ class Qwen2Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
dual_chunk_attention_config=dual_chunk_attention_config,
)
attn_cls = (
@@ -216,9 +213,7 @@ class Qwen2DecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 1000000)
- rope_scaling = getattr(config, "rope_scaling", None)
+ set_default_rope_theta(config, default_theta=1000000)
dual_chunk_attention_config = getattr(
config, "dual_chunk_attention_config", None
)
@@ -237,10 +232,9 @@ class Qwen2DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
cache_config=cache_config,
quant_config=quant_config,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
prefix=f"{prefix}.self_attn",
attn_type=attn_type,
dual_chunk_attention_config=dual_chunk_attention_config,
diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py
index 897dd7ef29f12..8c707c2561af1 100644
--- a/vllm/model_executor/models/qwen2_5_vl.py
+++ b/vllm/model_executor/models/qwen2_5_vl.py
@@ -46,7 +46,6 @@ from vllm.attention.layer import maybe_get_vit_flash_attn_backend
from vllm.attention.ops.vit_attn_wrappers import (
vit_flash_attn_wrapper,
vit_torch_sdpa_wrapper,
- vit_xformers_attn_wrapper,
)
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
@@ -64,6 +63,7 @@ from vllm.model_executor.layers.linear import (
RowParallelLinear,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.model_executor.models.vision import should_torch_compile_mm_vit
@@ -229,6 +229,9 @@ class Qwen2_5_VLVideoEmbeddingInputs(TensorSchema):
- hidden_size must match the hidden size of language model backbone.
- video_grid_thw shape: (num_videos, 3) in (grid_t, grid_h, grid_w)
format
+ - second_per_grid_ts: The video time interval (in seconds) for each
+ grid along the temporal dimension in the 3D position IDs. Returned
+ when `videos` is not `None`.
"""
type: Literal["video_embeds"]
@@ -243,6 +246,11 @@ class Qwen2_5_VLVideoEmbeddingInputs(TensorSchema):
TensorShape("nv", 3),
]
+ second_per_grid_ts: Annotated[
+ torch.Tensor | None,
+ TensorShape("nv"),
+ ] = None
+
Qwen2_5_VLVideoInputs: TypeAlias = (
Qwen2_5_VLVideoPixelInputs | Qwen2_5_VLVideoEmbeddingInputs
@@ -363,9 +371,9 @@ class Qwen2_5_VisionAttention(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: torch.Tensor, # Only used for Flash Attention
- seqlens: torch.Tensor, # Only used for xFormers
) -> torch.Tensor:
# [s, b, c] --> [s, b, head * 3 * head_dim]
x, _ = self.qkv(x)
@@ -378,13 +386,15 @@ class Qwen2_5_VisionAttention(nn.Module):
head=self.num_attention_heads_per_partition,
)
- if rotary_pos_emb is not None:
+ if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
qk, v = qkv[:, :, :2], qkv[:, :, 2]
qk_reshaped = einops.rearrange(
qk, "b s two head head_dim -> (two b) s head head_dim", two=2
)
- qk_rotated = apply_rotary_pos_emb_vision(qk_reshaped, rotary_pos_emb)
+ qk_rotated = apply_rotary_pos_emb_vision(
+ qk_reshaped, cos=rotary_pos_emb_cos, sin=rotary_pos_emb_sin
+ )
qk_rotated = qk_rotated.view(
2,
batch_size,
@@ -423,8 +433,6 @@ class Qwen2_5_VisionAttention(nn.Module):
v,
cu_seqlens,
)
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
output, _ = self.proj(context_layer)
return output
@@ -434,10 +442,9 @@ class Qwen2_5_VisionAttention(nn.Module):
dynamic_arg_dims={
"x": 0,
"cu_seqlens": 0,
- "rotary_pos_emb": 0,
- "seqlens": 0,
+ "rotary_pos_emb_cos": 0,
+ "rotary_pos_emb_sin": 0,
},
- mark_unbacked_dims={"seqlens": 0},
enable_if=should_torch_compile_mm_vit,
)
class Qwen2_5_VisionBlock(nn.Module):
@@ -485,16 +492,16 @@ class Qwen2_5_VisionBlock(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: torch.Tensor, # Only used for Flash Attention
- seqlens: torch.Tensor, # Only used for xFormers
) -> torch.Tensor:
x_attn = self.attn(
self.norm1(x),
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
x_fused_norm, residual = self.norm2(x, residual=x_attn)
x = residual + self.mlp(x_fused_norm)
@@ -588,42 +595,6 @@ class Qwen2_5_VisionPatchMerger(nn.Module):
return out
-class Qwen2_5_VisionRotaryEmbedding(nn.Module):
- def __init__(self, dim: int, theta: float = 10000.0) -> None:
- super().__init__()
- self.dim = dim
- self.theta = theta
- inv_freq = 1.0 / (
- theta ** (torch.arange(0, dim, 2, dtype=torch.float, device="cpu") / dim)
- )
- self.register_buffer("inv_freq", inv_freq, persistent=False)
- self._seq_len_cached = 0
- self._freqs_cached = None
-
- def update_freqs_cache(self, seqlen: int) -> None:
- if seqlen > self._seq_len_cached:
- seqlen *= 2
- self._seq_len_cached = seqlen
- self.inv_freq = 1.0 / (
- self.theta
- ** (
- torch.arange(
- 0, self.dim, 2, dtype=torch.float, device=self.inv_freq.device
- )
- / self.dim
- )
- )
- seq = torch.arange(
- seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
- )
- freqs = torch.outer(seq, self.inv_freq)
- self._freqs_cached = freqs
-
- def forward(self, seqlen: int) -> torch.Tensor:
- self.update_freqs_cache(seqlen)
- return self._freqs_cached[:seqlen]
-
-
class Qwen2_5_VisionTransformer(nn.Module):
def __init__(
self,
@@ -666,7 +637,12 @@ class Qwen2_5_VisionTransformer(nn.Module):
norm_layer = partial(RMSNorm, eps=norm_eps)
head_dim = self.hidden_size // self.num_heads
- self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+ self.rotary_pos_emb = get_rope(
+ head_size=head_dim,
+ rotary_dim=head_dim // 2,
+ max_position=8192,
+ is_neox_style=True,
+ )
use_upstream_fa = False
self.attn_backend = get_vit_attn_backend(
@@ -686,7 +662,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -757,15 +732,25 @@ class Qwen2_5_VisionTransformer(nn.Module):
)
pos_ids = torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)
max_size = max(h, w)
- rotary_pos_emb_full = self.rotary_pos_emb(max_size)
- rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
- rotary_pos_emb = rotary_pos_emb.reshape(
- rotary_pos_emb.shape[0] // self.spatial_merge_unit,
+
+ # Use pre-computed cos_sin_cache from RotaryEmbedding
+ cos, sin = self.rotary_pos_emb.get_cos_sin(max_size)
+
+ cos_combined = cos[pos_ids].flatten(1)
+ sin_combined = sin[pos_ids].flatten(1)
+
+ cos_combined = cos_combined.reshape(
+ cos_combined.shape[0] // self.spatial_merge_unit,
+ self.spatial_merge_unit,
+ -1,
+ )
+ sin_combined = sin_combined.reshape(
+ sin_combined.shape[0] // self.spatial_merge_unit,
self.spatial_merge_unit,
-1,
)
- return rotary_pos_emb
+ return cos_combined, sin_combined
def get_window_index_thw(self, grid_t, grid_h, grid_w):
vit_merger_window_size = (
@@ -807,14 +792,19 @@ class Qwen2_5_VisionTransformer(nn.Module):
@lru_cache(maxsize=1024) # noqa: B019
def get_rope_by_thw(self, t, h, w):
window_index_thw, cu_seqlens_window_thw = self.get_window_index_thw(t, h, w)
- rotary_pos_emb_thw = self.rotary_pos_emb_thw(t, h, w)
- rotary_pos_emb_thw = rotary_pos_emb_thw[window_index_thw, :, :]
- rotary_pos_emb_thw = rotary_pos_emb_thw.flatten(start_dim=0, end_dim=1)
+ cos_thw, sin_thw = self.rotary_pos_emb_thw(t, h, w)
+
+ cos_thw = cos_thw[window_index_thw, :, :]
+ cos_thw = cos_thw.flatten(start_dim=0, end_dim=1)
+ sin_thw = sin_thw[window_index_thw, :, :]
+ sin_thw = sin_thw.flatten(start_dim=0, end_dim=1)
+
cu_seqlens_thw = torch.repeat_interleave(
torch.tensor([h * w], dtype=torch.int32), t
)
return (
- rotary_pos_emb_thw,
+ cos_thw,
+ sin_thw,
window_index_thw,
cu_seqlens_window_thw,
cu_seqlens_thw,
@@ -823,17 +813,14 @@ class Qwen2_5_VisionTransformer(nn.Module):
def compute_attn_mask_seqlen(
self,
cu_seqlens: torch.Tensor,
- ) -> tuple[torch.Tensor, torch.Tensor]:
+ ) -> torch.Tensor:
max_seqlen = torch.zeros([], device=cu_seqlens.device)
- seqlens = torch.zeros(1, device=cu_seqlens.device)
if self.attn_backend in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.ROCM_AITER_FA,
}:
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
- return max_seqlen, seqlens
+ return max_seqlen
@staticmethod
def invert_permutation(perm: torch.Tensor) -> torch.Tensor:
@@ -849,7 +836,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
) -> torch.Tensor:
# patchify
seq_len, _ = x.size()
- rotary_pos_emb = []
+ rotary_pos_emb_cos = []
+ rotary_pos_emb_sin = []
window_index: list = []
cu_window_seqlens: list = [torch.tensor([0], dtype=torch.int32)]
cu_seqlens: list = []
@@ -865,7 +853,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
llm_w = w // self.spatial_merge_size
(
- rotary_pos_emb_thw,
+ cos_thw,
+ sin_thw,
window_index_thw,
cu_seqlens_window_thw,
cu_seqlens_thw,
@@ -878,11 +867,13 @@ class Qwen2_5_VisionTransformer(nn.Module):
cu_window_seqlens_last = cu_seqlens_window_thw[-1]
cu_window_seqlens.append(cu_seqlens_window_thw)
- rotary_pos_emb.append(rotary_pos_emb_thw)
+ rotary_pos_emb_cos.append(cos_thw)
+ rotary_pos_emb_sin.append(sin_thw)
cu_seqlens.append(cu_seqlens_thw)
- rotary_pos_emb = torch.cat(rotary_pos_emb)
+ rotary_pos_emb_cos = torch.cat(rotary_pos_emb_cos)
+ rotary_pos_emb_sin = torch.cat(rotary_pos_emb_sin)
window_index = torch.cat(window_index)
# compute reverse indices
reverse_indices = self.invert_permutation(window_index)
@@ -894,14 +885,17 @@ class Qwen2_5_VisionTransformer(nn.Module):
# transformers
# pre-compute seqlens for window/full attn to reduce cuMemcpy operations
- max_seqlen_full, seqlens_full = self.compute_attn_mask_seqlen(cu_seqlens)
- max_seqlen_window, seqlens_window = self.compute_attn_mask_seqlen(
- cu_window_seqlens
- )
+ max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens)
+ max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens)
cu_seqlens = cu_seqlens.to(device=self.device, non_blocking=True)
cu_window_seqlens = cu_window_seqlens.to(device=self.device, non_blocking=True)
- rotary_pos_emb = rotary_pos_emb.to(device=self.device, non_blocking=True)
+ rotary_pos_emb_cos = rotary_pos_emb_cos.to(
+ device=self.device, non_blocking=True
+ )
+ rotary_pos_emb_sin = rotary_pos_emb_sin.to(
+ device=self.device, non_blocking=True
+ )
window_index = window_index.to(device=hidden_states.device, non_blocking=True)
reverse_indices = reverse_indices.to(
device=hidden_states.device, non_blocking=True
@@ -919,18 +913,16 @@ class Qwen2_5_VisionTransformer(nn.Module):
if layer_num in self.fullatt_block_indexes:
cu_seqlens_now = cu_seqlens
max_seqlen_now = max_seqlen_full
- seqlens_now = seqlens_full
else:
cu_seqlens_now = cu_window_seqlens
max_seqlen_now = max_seqlen_window
- seqlens_now = seqlens_window
hidden_states = blk(
hidden_states,
cu_seqlens=cu_seqlens_now,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen_now,
- seqlens=seqlens_now,
)
# For Qwen2.5-VL-3B, float16 will overflow at last block
@@ -1310,6 +1302,7 @@ class Qwen2_5_VLForConditionalGeneration(
type="video_embeds",
video_embeds=video_embeds,
video_grid_thw=video_grid_thw,
+ second_per_grid_ts=second_per_grid_ts,
)
def _process_image_input(
@@ -1421,7 +1414,13 @@ class Qwen2_5_VLForConditionalGeneration(
# Cast to long to match the original code
# https://github.com/huggingface/transformers/blob/41980ce93e775f6c88500c51c8db7946fc6a2add/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py#L491 # noqa
- second_per_grid_ts = video_input["second_per_grid_ts"].long()
+ second_per_grid_ts = video_input.get("second_per_grid_ts")
+ if second_per_grid_ts is None:
+ raise ValueError(
+ "second_per_grid_ts is required when video_pruning_rate > 0 "
+ "is enabled for video inputs, including the video_embeds path."
+ )
+ second_per_grid_ts = second_per_grid_ts.long()
tokens_per_second = self.config.vision_config.tokens_per_second
video_embeds_out = []
diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py
index 2ff0d19df238c..6b97d0b2ca2e3 100644
--- a/vllm/model_executor/models/qwen2_moe.py
+++ b/vllm/model_executor/models/qwen2_moe.py
@@ -194,8 +194,7 @@ class Qwen2MoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
@@ -222,7 +221,6 @@ class Qwen2MoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.dual_chunk_attention_config = dual_chunk_attention_config
@@ -248,8 +246,7 @@ class Qwen2MoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
dual_chunk_attention_config=dual_chunk_attention_config,
)
self.attn = Attention(
@@ -291,8 +288,6 @@ class Qwen2MoeDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
dual_chunk_attention_config = getattr(
config, "dual_chunk_attention_config", None
)
@@ -301,8 +296,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
cache_config=cache_config,
quant_config=quant_config,
diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py
index 5d21e249fc4cc..9d1d023aed172 100644
--- a/vllm/model_executor/models/qwen2_vl.py
+++ b/vllm/model_executor/models/qwen2_vl.py
@@ -29,10 +29,11 @@ from collections.abc import Callable, Iterable, Mapping, Sequence
from functools import partial
from typing import Annotated, Any, Literal, TypeAlias
+import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
-from einops import rearrange, repeat
+from einops import rearrange
from transformers import BatchFeature
from transformers.models.qwen2_vl import Qwen2VLImageProcessor, Qwen2VLProcessor
from transformers.models.qwen2_vl.configuration_qwen2_vl import (
@@ -59,7 +60,9 @@ from vllm.model_executor.layers.linear import (
RowParallelLinear,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.rotary_embedding.common import (
+ apply_rotary_emb_torch,
dispatch_rotary_emb_function,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -275,47 +278,13 @@ class Qwen2VisionMLP(nn.Module):
return x
-def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
- if not interleaved:
- x1, x2 = x.chunk(2, dim=-1)
- return torch.cat((-x2, x1), dim=-1)
- else:
- x1, x2 = x[..., ::2], x[..., 1::2]
- return rearrange(
- torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2
- )
-
-
-def apply_rotary_emb_torch(
- x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
+def apply_rotary_pos_emb_vision(
+ t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
- """
- x: (batch_size, seqlen, nheads, headdim)
- cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
- """
- ro_dim = cos.shape[-1] * 2
- assert ro_dim <= x.shape[-1]
- cos = repeat(
- cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
+ rotary_emb_function = dispatch_rotary_emb_function(
+ default=partial(apply_rotary_emb_torch, is_neox_style=True)
)
- sin = repeat(
- sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
- )
- return torch.cat(
- [
- x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin,
- x[..., ro_dim:],
- ],
- dim=-1,
- )
-
-
-def apply_rotary_pos_emb_vision(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
- rotary_emb_function = dispatch_rotary_emb_function(default=apply_rotary_emb_torch)
- t_ = t.float()
- cos = freqs.cos()
- sin = freqs.sin()
- output = rotary_emb_function(t_, cos, sin).type_as(t)
+ output = rotary_emb_function(t, cos, sin).type_as(t)
return output
@@ -379,7 +348,6 @@ class Qwen2VisionAttention(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -412,9 +380,9 @@ class Qwen2VisionAttention(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: int | None = None, # Only used for Flash Attention
- seqlens: list[int] | None = None, # Only used for xFormers
) -> torch.Tensor:
# [s, b, c] --> [s, b, 3 * head * head_dim]
x, _ = self.qkv(x)
@@ -424,11 +392,13 @@ class Qwen2VisionAttention(nn.Module):
batch_size = q.shape[1]
q, k, v = (rearrange(x, "s b ... -> b s ...") for x in (q, k, v))
- if rotary_pos_emb is not None:
- # [2 * b, s, heads, head_dim]
- qk_concat = torch.cat([q, k], dim=0)
- qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
- q, k = torch.chunk(qk_rotated, 2, dim=0)
+
+ # [2 * b, s, heads, head_dim]
+ qk_concat = torch.cat([q, k], dim=0)
+ qk_rotated = apply_rotary_pos_emb_vision(
+ qk_concat, rotary_pos_emb_cos, rotary_pos_emb_sin
+ )
+ q, k = torch.chunk(qk_rotated, 2, dim=0)
if self.is_flash_attn_backend:
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
@@ -473,20 +443,6 @@ class Qwen2VisionAttention(nn.Module):
context_layer = rearrange(
context_layer, "b s h d -> s b (h d)"
).contiguous()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import BlockDiagonalMask
-
- attn_bias = BlockDiagonalMask.from_seqlens(
- q_seqlen=seqlens, kv_seqlen=None, device=q.device
- )
-
- context_layer = xops.memory_efficient_attention_forward(
- q, k, v, attn_bias=attn_bias, p=0, scale=None
- )
- context_layer = rearrange(
- context_layer, "b s h d -> s b (h d)"
- ).contiguous()
output, _ = self.proj(context_layer)
return output
@@ -534,16 +490,16 @@ class Qwen2VisionBlock(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: int | None = None, # Only used for Flash Attention
- seqlens: list[int] | None = None, # Only used for xFormers
) -> torch.Tensor:
x = x + self.attn(
self.norm1(x),
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
x = x + self.mlp(self.norm2(x))
@@ -628,40 +584,6 @@ class Qwen2VisionPatchMerger(nn.Module):
return out
-class Qwen2VisionRotaryEmbedding(nn.Module):
- def __init__(self, dim: int, theta: float = 10000.0) -> None:
- super().__init__()
- self.dim = dim
- self.theta = theta
- inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
- self._seq_len_cached = 0
- self._freqs_cached = None
-
- def update_freqs_cache(self, seqlen: int) -> None:
- if seqlen > self._seq_len_cached:
- seqlen *= 2
- self._seq_len_cached = seqlen
- self.inv_freq = 1.0 / (
- self.theta
- ** (
- torch.arange(
- 0, self.dim, 2, dtype=torch.float, device=self.inv_freq.device
- )
- / self.dim
- )
- )
- seq = torch.arange(
- seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
- )
- freqs = torch.outer(seq, self.inv_freq)
- self._freqs_cached = freqs
-
- def forward(self, seqlen: int) -> torch.Tensor:
- self.update_freqs_cache(seqlen)
- return self._freqs_cached[:seqlen]
-
-
class Qwen2VisionTransformer(nn.Module):
def __init__(
self,
@@ -700,7 +622,12 @@ class Qwen2VisionTransformer(nn.Module):
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
head_dim = embed_dim // num_heads
- self.rotary_pos_emb = Qwen2VisionRotaryEmbedding(head_dim // 2)
+ self.rotary_pos_emb = get_rope(
+ head_size=head_dim,
+ rotary_dim=head_dim // 2,
+ max_position=8192,
+ is_neox_style=True,
+ )
self.blocks = nn.ModuleList(
[
@@ -744,7 +671,9 @@ class Qwen2VisionTransformer(nn.Module):
def device(self) -> torch.device:
return self.patch_embed.proj.weight.device
- def rot_pos_emb(self, grid_thw: list[list[int]]) -> torch.Tensor:
+ def rot_pos_emb(
+ self, grid_thw: list[list[int]]
+ ) -> tuple[torch.Tensor, torch.Tensor]:
pos_ids = []
max_grid_size = 0
for t, h, w in grid_thw:
@@ -773,22 +702,22 @@ class Qwen2VisionTransformer(nn.Module):
pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
max_grid_size = max(max_grid_size, h, w)
pos_ids = torch.cat(pos_ids, dim=0)
- rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
- rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
- return rotary_pos_emb
- def compute_attn_mask_seqlen(
- self, cu_seqlens: torch.Tensor
- ) -> tuple[int | None, list[int] | None]:
- max_seqlen, seqlens = None, None
+ # Use pre-computed cos_sin_cache from RotaryEmbedding
+ cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)
+
+ cos_combined = cos[pos_ids].flatten(1)
+ sin_combined = sin[pos_ids].flatten(1)
+ return cos_combined, sin_combined
+
+ def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
+ max_seqlen = None
if self.attn_backend in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.ROCM_AITER_FA,
}:
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
- return max_seqlen, seqlens
+ return max_seqlen
def forward(
self,
@@ -801,32 +730,34 @@ class Qwen2VisionTransformer(nn.Module):
if isinstance(grid_thw, list):
grid_thw_list = grid_thw
- grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
+ grid_thw = np.array(grid_thw, dtype=np.int32)
else:
grid_thw_list = grid_thw.tolist()
+ grid_thw = grid_thw.numpy()
# compute position embedding
- rotary_pos_emb = self.rot_pos_emb(grid_thw_list)
+ rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
# compute cu_seqlens
- cu_seqlens = torch.repeat_interleave(
- grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
- ).cumsum(dim=0, dtype=torch.int32)
- cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
- cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
+ cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ axis=0, dtype=np.int32
+ )
+ cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
+ cu_seqlens = torch.from_numpy(cu_seqlens)
# transformers
x = x.unsqueeze(1)
# pre-compute seqlens for attn mask to reduce cuMemcpy operations
- max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
+ cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
for blk in self.blocks:
x = blk(
x,
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
# adapter
diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py
index 8d7f22a33fe6c..93a629d81e8ff 100644
--- a/vllm/model_executor/models/qwen3.py
+++ b/vllm/model_executor/models/qwen3.py
@@ -42,6 +42,7 @@ 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.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
from .qwen2 import Qwen2MLP as Qwen3MLP
@@ -57,14 +58,13 @@ class Qwen3Attention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
+ rope_parameters: dict,
max_position: int = 4096 * 32,
head_dim: int | None = None,
rms_norm_eps: float = 1e-06,
qkv_bias: bool = False,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
- rope_scaling: tuple | None = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER,
dual_chunk_attention_config: dict[str, Any] | None = None,
@@ -89,7 +89,6 @@ class Qwen3Attention(nn.Module):
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.rope_theta = rope_theta
self.dual_chunk_attention_config = dual_chunk_attention_config
self.qkv_proj = QKVParallelLinear(
@@ -113,8 +112,7 @@ class Qwen3Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
dual_chunk_attention_config=dual_chunk_attention_config,
)
self.attn = Attention(
@@ -166,9 +164,7 @@ class Qwen3DecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 1000000)
- rope_scaling = getattr(config, "rope_scaling", None)
+ set_default_rope_theta(config, default_theta=1000000)
dual_chunk_attention_config = getattr(
config, "dual_chunk_attention_config", None
)
@@ -187,13 +183,12 @@ class Qwen3DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
rms_norm_eps=config.rms_norm_eps,
qkv_bias=getattr(config, "attention_bias", False),
head_dim=getattr(config, "head_dim", None),
cache_config=cache_config,
quant_config=quant_config,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
prefix=f"{prefix}.self_attn",
attn_type=attn_type,
dual_chunk_attention_config=dual_chunk_attention_config,
diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py
index 96751fee800bb..8ee3dd99e11db 100644
--- a/vllm/model_executor/models/qwen3_moe.py
+++ b/vllm/model_executor/models/qwen3_moe.py
@@ -216,8 +216,7 @@ class Qwen3MoeAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any],
max_position_embeddings: int = 8192,
head_dim: int | None = None,
rms_norm_eps: float = 1e-06,
@@ -247,7 +246,6 @@ class Qwen3MoeAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.dual_chunk_attention_config = dual_chunk_attention_config
@@ -273,8 +271,7 @@ class Qwen3MoeAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
dual_chunk_attention_config=dual_chunk_attention_config,
)
self.attn = Attention(
@@ -326,8 +323,6 @@ class Qwen3MoeDecoderLayer(nn.Module):
quant_config = vllm_config.quant_config
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
dual_chunk_attention_config = getattr(
config, "dual_chunk_attention_config", None
@@ -336,8 +331,7 @@ class Qwen3MoeDecoderLayer(nn.Module):
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
max_position_embeddings=max_position_embeddings,
rms_norm_eps=config.rms_norm_eps,
qkv_bias=getattr(config, "attention_bias", False),
diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py
index 86508a7c64317..bfed64728305e 100644
--- a/vllm/model_executor/models/qwen3_next.py
+++ b/vllm/model_executor/models/qwen3_next.py
@@ -10,7 +10,7 @@ from einops import rearrange
from torch import nn
from transformers.activations import ACT2FN
-from vllm.attention import Attention, AttentionBackend, AttentionMetadata
+from vllm.attention import Attention, AttentionMetadata
from vllm.compilation.decorators import support_torch_compile
from vllm.config import (
CacheConfig,
@@ -216,12 +216,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
@property
def mamba_type(self) -> str:
- return "linear_attention"
-
- def get_attn_backend(self) -> type["AttentionBackend"]:
- from vllm.v1.attention.backends.gdn_attn import GDNAttentionBackend
-
- return GDNAttentionBackend
+ return "gdn_attention"
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
@@ -753,8 +748,7 @@ class Qwen3NextAttention(nn.Module):
head_size=self.head_dim,
rotary_dim=self.head_dim,
max_position=config.max_position_embeddings,
- base=config.rope_theta,
- rope_scaling=config.rope_scaling,
+ rope_parameters=config.rope_parameters,
partial_rotary_factor=config.partial_rotary_factor,
dual_chunk_attention_config=self.dual_chunk_attention_config,
)
@@ -1154,8 +1148,8 @@ class QwenNextMixtureOfExperts(MixtureOfExperts):
example_moe = layer.mlp
self.moe_layers.append(layer.mlp.experts)
- if example_moe is None:
- raise RuntimeError("No Qwen3Next layer found in the model.layers.")
+ if example_moe is None:
+ raise RuntimeError("No Qwen3Next layer found in the model.layers.")
# Set MoE hyperparameters
self.num_moe_layers = len(self.moe_layers)
diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py
index 40b80ce2387c8..61f218f16d79c 100755
--- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py
+++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py
@@ -60,6 +60,7 @@ 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.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
from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo
@@ -90,7 +91,6 @@ from .qwen2_5_omni_thinker import (
)
from .qwen2_5_vl import (
Qwen2_5_VisionAttention,
- Qwen2_5_VisionRotaryEmbedding,
Qwen2_5_VLProcessingInfo,
)
from .qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel
@@ -221,16 +221,16 @@ class Qwen3_VisionBlock(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: torch.Tensor, # Only used for Flash Attention
- seqlens: torch.Tensor, # Only used for xFormers
) -> torch.Tensor:
x = x + self.attn(
self.norm1(x),
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
x = x + self.mlp(self.norm2(x))
@@ -332,7 +332,12 @@ class Qwen3Omni_VisionTransformer(nn.Module):
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
head_dim = self.hidden_size // self.num_heads
- self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+ self.rotary_pos_emb = get_rope(
+ head_size=head_dim,
+ rotary_dim=head_dim // 2,
+ max_position=8192,
+ is_neox_style=True,
+ )
self.blocks = nn.ModuleList(
[
@@ -416,9 +421,14 @@ class Qwen3Omni_VisionTransformer(nn.Module):
pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
pos_ids = torch.cat(pos_ids, dim=0)
max_grid_size = grid_thw[:, 1:].max()
- rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
- rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
- return rotary_pos_emb
+
+ # Use pre-computed cos_sin_cache from RotaryEmbedding
+ cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)
+
+ cos_combined = cos[pos_ids].flatten(1)
+ sin_combined = sin[pos_ids].flatten(1)
+
+ return cos_combined, sin_combined
def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
num_grid_per_side = self.num_grid_per_side
@@ -488,14 +498,11 @@ class Qwen3Omni_VisionTransformer(nn.Module):
def compute_attn_mask_seqlen(
self,
cu_seqlens: torch.Tensor,
- ) -> tuple[torch.Tensor, torch.Tensor]:
+ ) -> torch.Tensor:
max_seqlen = torch.zeros([], device=cu_seqlens.device)
- seqlens = torch.zeros(1, device=cu_seqlens.device)
if self.attn_backend == AttentionBackendEnum.FLASH_ATTN:
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
- return max_seqlen, seqlens
+ return max_seqlen
def forward(
self,
@@ -508,7 +515,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):
if self.apply_vit_abs_pos_embed:
pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
hidden_states = hidden_states + pos_embeds
- rotary_pos_emb = self.rot_pos_emb(grid_thw)
+ rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw)
cu_seqlens = torch.repeat_interleave(
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
@@ -519,8 +526,9 @@ class Qwen3Omni_VisionTransformer(nn.Module):
cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
hidden_states = hidden_states.unsqueeze(1)
- rotary_pos_emb = rotary_pos_emb.to(hidden_states.device)
- max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+ rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device)
+ rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device)
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
hidden_states_list = []
deepstack_visual_indexes = self.deepstack_visual_indexes
@@ -529,9 +537,9 @@ class Qwen3Omni_VisionTransformer(nn.Module):
hidden_states = blk(
hidden_states,
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
if (
deepstack_visual_indexes is not None
diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py
index 7f0c9372991d1..4cd6fa14c32df 100644
--- a/vllm/model_executor/models/qwen3_vl.py
+++ b/vllm/model_executor/models/qwen3_vl.py
@@ -63,6 +63,7 @@ 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.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
from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -95,7 +96,6 @@ from .interfaces import (
)
from .qwen2_5_vl import (
Qwen2_5_VisionAttention,
- Qwen2_5_VisionRotaryEmbedding,
Qwen2_5_VLImageEmbeddingInputs,
Qwen2_5_VLImageInputs,
Qwen2_5_VLImagePixelInputs,
@@ -232,16 +232,16 @@ class Qwen3_VisionBlock(nn.Module):
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
- rotary_pos_emb: torch.Tensor,
+ rotary_pos_emb_cos: torch.Tensor,
+ rotary_pos_emb_sin: torch.Tensor,
max_seqlen: torch.Tensor, # Only used for Flash Attention
- seqlens: torch.Tensor, # Only used for xFormers
) -> torch.Tensor:
x = x + self.attn(
self.norm1(x),
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
x = x + self.mlp(self.norm2(x))
@@ -339,7 +339,12 @@ class Qwen3_VisionTransformer(nn.Module):
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
head_dim = self.hidden_size // self.num_heads
- self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+ self.rotary_pos_emb = get_rope(
+ head_size=head_dim,
+ rotary_dim=head_dim // 2,
+ max_position=8192,
+ is_neox_style=True,
+ )
self.merger = Qwen3_VisionPatchMerger(
d_model=vision_config.out_hidden_size,
@@ -384,7 +389,6 @@ class Qwen3_VisionTransformer(nn.Module):
if self.attn_backend not in {
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
- AttentionBackendEnum.XFORMERS,
AttentionBackendEnum.ROCM_AITER_FA,
}:
raise RuntimeError(
@@ -451,10 +455,15 @@ class Qwen3_VisionTransformer(nn.Module):
else self.rot_pos_ids(h, w, self.spatial_merge_size).repeat(t, 1)
for t, h, w in grid_thw
]
- pos_ids = torch.cat(pos_ids, dim=0)
- rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
- rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
- return rotary_pos_emb
+ pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True)
+
+ # Use pre-computed cos_sin_cache from RotaryEmbedding
+ cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)
+
+ cos_combined = cos[pos_ids].flatten(1)
+ sin_combined = sin[pos_ids].flatten(1)
+
+ return cos_combined, sin_combined
def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
num_grid_per_side = self.num_grid_per_side
@@ -519,17 +528,14 @@ class Qwen3_VisionTransformer(nn.Module):
def compute_attn_mask_seqlen(
self,
cu_seqlens: torch.Tensor,
- ) -> tuple[torch.Tensor, torch.Tensor]:
+ ) -> torch.Tensor:
max_seqlen = torch.zeros([], device=cu_seqlens.device)
- seqlens = torch.zeros(1, device=cu_seqlens.device)
if (
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
):
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
- elif self.attn_backend == AttentionBackendEnum.XFORMERS:
- seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
- return max_seqlen, seqlens
+ return max_seqlen
def forward(
self,
@@ -541,22 +547,23 @@ class Qwen3_VisionTransformer(nn.Module):
if isinstance(grid_thw, list):
grid_thw_list = grid_thw
- grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
+ grid_thw = np.array(grid_thw, dtype=np.int32)
else:
grid_thw_list = grid_thw.tolist()
+ grid_thw = grid_thw.numpy()
pos_embeds = self.fast_pos_embed_interpolate(grid_thw_list)
hidden_states = hidden_states + pos_embeds
- rotary_pos_emb = self.rot_pos_emb(grid_thw_list)
- rotary_pos_emb = rotary_pos_emb.to(hidden_states.device, non_blocking=True)
+ rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
- cu_seqlens = torch.repeat_interleave(
- grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
- ).cumsum(dim=0, dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32)
- cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
+ cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ axis=0, dtype=np.int32
+ )
+ cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
+ cu_seqlens = torch.from_numpy(cu_seqlens)
hidden_states = hidden_states.unsqueeze(1)
- max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
deepstack_feature_lists = []
@@ -564,9 +571,9 @@ class Qwen3_VisionTransformer(nn.Module):
hidden_states = blk(
hidden_states,
cu_seqlens=cu_seqlens,
- rotary_pos_emb=rotary_pos_emb,
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
max_seqlen=max_seqlen,
- seqlens=seqlens,
)
if layer_num in self.deepstack_visual_indexes:
deepstack_merger_idx = self.deepstack_visual_indexes.index(layer_num)
diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py
index 5c3205faf9c2f..e2c129120b1a5 100644
--- a/vllm/model_executor/models/qwen3_vl_moe.py
+++ b/vllm/model_executor/models/qwen3_vl_moe.py
@@ -15,7 +15,7 @@
# 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
+# 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,
@@ -29,7 +29,9 @@ from collections.abc import Callable, Iterable
from itertools import islice
import torch
-from transformers.models.qwen3_vl_moe.configuration_qwen3_vl_moe import Qwen3VLMoeConfig
+from transformers.models.qwen3_vl_moe.configuration_qwen3_vl_moe import (
+ Qwen3VLMoeConfig,
+)
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
@@ -44,7 +46,12 @@ from vllm.model_executor.model_loader.weight_utils import (
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.sequence import IntermediateTensors
-from .qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel
+from .interfaces import MixtureOfExperts
+from .qwen3_moe import (
+ Qwen3MoeForCausalLM,
+ Qwen3MoeModel,
+ Qwen3MoeSparseMoeBlock,
+)
from .qwen3_vl import (
Qwen3_VisionTransformer,
Qwen3VLDummyInputsBuilder,
@@ -344,12 +351,56 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):
)
+class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts):
+ def update_physical_experts_metadata(
+ self,
+ num_physical_experts: int,
+ num_local_physical_experts: int,
+ ) -> None:
+ 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.language_model.model.layers:
+ if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
+ 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 set_moe_parameters(self):
+ self.expert_weights = []
+
+ self.moe_layers = []
+ example_moe = None
+ for layer in self.language_model.model.layers:
+ if hasattr(layer, "mlp") and isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
+ example_moe = layer.mlp
+ self.moe_layers.append(layer.mlp.experts)
+
+ if example_moe is None:
+ raise RuntimeError("No Qwen3Moe layer found in the language_model.")
+
+ # Set MoE hyperparameters
+ self.num_moe_layers = len(self.moe_layers)
+ self.num_expert_groups = 1
+ self.num_shared_experts = 0
+ self.num_logical_experts = example_moe.n_logical_experts
+ self.num_physical_experts = example_moe.n_physical_experts
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
+ self.num_routed_experts = example_moe.n_routed_experts
+ self.num_redundant_experts = example_moe.n_redundant_experts
+
+
@MULTIMODAL_REGISTRY.register_processor(
Qwen3VLMultiModalProcessor,
info=Qwen3VLMoeProcessingInfo,
dummy_inputs=Qwen3VLDummyInputsBuilder,
)
-class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
+class Qwen3VLMoeForConditionalGeneration(
+ Qwen3VLForConditionalGeneration, Qwen3VLMoeMixtureOfExperts
+):
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -413,3 +464,6 @@ class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
self.deepstack_input_embeds = None
self.visual_dim = config.vision_config.out_hidden_size
self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+ # Set MoE hyperparameters
+ self.set_moe_parameters()
diff --git a/vllm/model_executor/models/qwen_vl.py b/vllm/model_executor/models/qwen_vl.py
index 6a259cade9cf1..4906cf441f6fb 100644
--- a/vllm/model_executor/models/qwen_vl.py
+++ b/vllm/model_executor/models/qwen_vl.py
@@ -25,6 +25,7 @@ from transformers.tokenization_utils_base import TextInput
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
ReplicatedLinear,
@@ -333,7 +334,7 @@ class VisionTransformer(nn.Module):
patch_height, patch_width = self.patch_size = (patch_size, patch_size)
self.grid_size = (image_height // patch_height, image_width // patch_width)
self.output_dim = output_dim
- self.conv1 = nn.Conv2d(
+ self.conv1 = Conv2dLayer(
in_channels=3,
out_channels=width,
kernel_size=patch_size,
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index 6e9790de49bfa..b3da64af750c7 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -157,6 +157,7 @@ _TEXT_GENERATION_MODELS = {
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
"PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
"Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
+ "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
"QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
"Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
@@ -288,6 +289,10 @@ _MULTIMODAL_MODELS = {
"H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
"InternVLChatModel": ("internvl", "InternVLChatModel"),
"NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
+ "OpenCUAForConditionalGeneration": (
+ "opencua",
+ "OpenCUAForConditionalGeneration",
+ ),
"InternS1ForConditionalGeneration": (
"interns1",
"InternS1ForConditionalGeneration",
@@ -597,7 +602,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):
mi_dict = json.load(file)
except FileNotFoundError:
logger.debug(
- ("Cached model info file for class %s.%s not found"),
+ "Cached model info file for class %s.%s not found",
self.module_name,
self.class_name,
)
@@ -605,7 +610,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):
if mi_dict["hash"] != module_hash:
logger.debug(
- ("Cached model info file for class %s.%s is stale"),
+ "Cached model info file for class %s.%s is stale",
self.module_name,
self.class_name,
)
@@ -615,7 +620,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):
return _ModelInfo(**mi_dict["modelinfo"])
except Exception:
logger.debug(
- ("Cached model info for class %s.%s error. "),
+ "Cached model info for class %s.%s error. ",
self.module_name,
self.class_name,
)
@@ -650,14 +655,14 @@ class _LazyRegisteredModel(_BaseRegisteredModel):
mi = self._load_modelinfo_from_cache(module_hash)
if mi is not None:
logger.debug(
- ("Loaded model info for class %s.%s from cache"),
+ "Loaded model info for class %s.%s from cache",
self.module_name,
self.class_name,
)
return mi
else:
logger.debug(
- ("Cache model info for class %s.%s miss. Loading model instead."),
+ "Cache model info for class %s.%s miss. Loading model instead.",
self.module_name,
self.class_name,
)
diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py
index bf211d28f1844..4744d8e44f390 100644
--- a/vllm/model_executor/models/seed_oss.py
+++ b/vllm/model_executor/models/seed_oss.py
@@ -54,6 +54,7 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import set_default_rope_theta
from .interfaces import SupportsLoRA, SupportsPP
from .utils import (
@@ -112,11 +113,10 @@ class SeedOssAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
head_dim: int,
+ rope_parameters: dict,
max_position: int = 4096 * 32,
- rope_theta: float = 10000,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
- rope_scaling: tuple | None = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER,
) -> None:
@@ -140,7 +140,6 @@ class SeedOssAttention(nn.Module):
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.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
@@ -163,8 +162,7 @@ class SeedOssAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
- base=self.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -200,9 +198,7 @@ class SeedOssDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- # Requires transformers > 4.32.0
- rope_theta = getattr(config, "rope_theta", 1000000)
- rope_scaling = getattr(config, "rope_scaling", None)
+ set_default_rope_theta(config, default_theta=1000000)
# By default, SeedOss uses causal attention as it is a
# decoder-only model.
@@ -219,10 +215,9 @@ class SeedOssDecoderLayer(nn.Module):
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
head_dim=config.head_dim,
- rope_theta=rope_theta,
cache_config=cache_config,
quant_config=quant_config,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
prefix=f"{prefix}.self_attn",
attn_type=attn_type,
)
diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py
index b175dd60cf650..ce5847bf79a5e 100644
--- a/vllm/model_executor/models/siglip.py
+++ b/vllm/model_executor/models/siglip.py
@@ -24,6 +24,7 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -286,7 +287,7 @@ class SiglipVisionEmbeddings(nn.Module):
self.image_size = config.image_size
self.patch_size = config.patch_size
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=config.num_channels,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
@@ -827,6 +828,7 @@ class SiglipVisionModel(nn.Module):
) -> None:
super().__init__()
+ self.quant_config = quant_config
self.vision_model = SiglipVisionTransformer(
config,
quant_config,
@@ -911,12 +913,38 @@ class SiglipVisionModel(nn.Module):
break
else:
param = params_dict[name]
+ param = maybe_swap_ffn_param(
+ name, param, loaded_weight, params_dict, self.quant_config
+ )
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params
+def maybe_swap_ffn_param(
+ name: str,
+ param: torch.Tensor,
+ loaded_weight: torch.Tensor,
+ params_dict: dict[str, torch.Tensor],
+ quant_config: QuantizationConfig,
+) -> torch.Tensor:
+ if not (quant_config and quant_config.get_name() == "gguf") or ".fc" not in name:
+ return param
+ # Some GGUF models have fc1 and fc2 weights swapped
+ tp_size = get_tensor_model_parallel_world_size()
+ output_dim = getattr(param, "output_dim", 0)
+ output_size = param.size(output_dim) * tp_size
+ weight_out_size = loaded_weight.size(output_dim)
+ if ".fc1." in name and output_size != weight_out_size:
+ new_name = name.replace(".fc1.", ".fc2.")
+ param = params_dict[new_name]
+ elif ".fc2." in name and output_size != weight_out_size:
+ new_name = name.replace(".fc2.", ".fc1.")
+ param = params_dict[new_name]
+ return param
+
+
# Adapted from: https://github.com/huggingface/transformers/blob/v4.54.1/src/transformers/models/siglip/modeling_siglip.py#L200
class SiglipTextEmbeddings(nn.Module):
def __init__(self, config: SiglipTextConfig):
diff --git a/vllm/model_executor/models/siglip2navit.py b/vllm/model_executor/models/siglip2navit.py
index 29dd164ad37fd..c185b45345bd5 100644
--- a/vllm/model_executor/models/siglip2navit.py
+++ b/vllm/model_executor/models/siglip2navit.py
@@ -16,6 +16,7 @@ from vllm.attention.backends.registry import AttentionBackendEnum
from vllm.attention.layer import maybe_get_vit_flash_attn_backend
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
LinearBase,
@@ -67,7 +68,7 @@ class Siglip2VisionEmbeddings(nn.Module):
self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
else:
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=config.num_channels,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
@@ -99,7 +100,7 @@ class Siglip2VisionEmbeddings(nn.Module):
target_dtype = self.patch_embedding.weight.dtype
if isinstance(self.patch_embedding, LinearBase):
patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
- elif isinstance(self.patch_embedding, nn.Conv2d):
+ elif isinstance(self.patch_embedding, Conv2dLayer):
pixel_values = pixel_values.view(
-1,
self.config.num_channels * self.config.temporal_patch_size,
@@ -190,7 +191,7 @@ def apply_rotary_pos_emb(
cos = cos.chunk(2, dim=-1)[0].contiguous()
sin = sin.chunk(2, dim=-1)[0].contiguous()
if is_flash_attn_backend and not current_platform.is_xpu():
- from flash_attn.layers.rotary import apply_rotary_emb
+ from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
apply_rotary_emb_func = apply_rotary_emb
else:
diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py
index 4ec855f794446..7e9fc51036d2e 100644
--- a/vllm/model_executor/models/solar.py
+++ b/vllm/model_executor/models/solar.py
@@ -25,7 +25,6 @@
"""Inference-only Solar model compatible with HuggingFace weights."""
from collections.abc import Iterable
-from typing import Any
import torch
from torch import nn
@@ -111,8 +110,6 @@ class SolarAttention(nn.Module):
hidden_size: int,
num_heads: int,
num_kv_heads: int,
- rope_theta: float = 10000,
- rope_scaling: dict[str, Any] | None = None,
max_position_embeddings: int = 8192,
quant_config: QuantizationConfig | None = None,
bias: bool = False,
@@ -142,7 +139,6 @@ class SolarAttention(nn.Module):
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.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.qkv_proj = QKVParallelLinear(
@@ -166,8 +162,7 @@ class SolarAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embeddings,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
)
self.attn = Attention(
self.num_heads,
@@ -202,15 +197,6 @@ class SolarDecoderLayer(nn.Module):
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
- rope_theta = getattr(config, "rope_theta", 10000)
- rope_scaling = getattr(config, "rope_scaling", None)
-
- if rope_scaling is not None and getattr(
- config, "original_max_position_embeddings", None
- ):
- rope_scaling["original_max_position_embeddings"] = (
- config.original_max_position_embeddings
- )
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Support abacusai/Smaug-72B-v0.1 with attention_bias
# Support internlm/internlm-7b with bias
@@ -224,8 +210,6 @@ class SolarDecoderLayer(nn.Module):
num_kv_heads=getattr(
config, "num_key_value_heads", config.num_attention_heads
),
- rope_theta=rope_theta,
- rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
quant_config=quant_config,
bias=attention_bias,
diff --git a/vllm/model_executor/models/stablelm.py b/vllm/model_executor/models/stablelm.py
index 06eb7201c1a89..a738fcbb4ee28 100644
--- a/vllm/model_executor/models/stablelm.py
+++ b/vllm/model_executor/models/stablelm.py
@@ -153,7 +153,7 @@ class StablelmAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.config.max_position_embeddings,
- base=self.config.rope_theta,
+ rope_parameters=self.config.rope_parameters,
partial_rotary_factor=self.partial_rotary_factor,
)
self.attn = Attention(
diff --git a/vllm/model_executor/models/starcoder2.py b/vllm/model_executor/models/starcoder2.py
index 0f2942acd5006..1118fca3cac91 100644
--- a/vllm/model_executor/models/starcoder2.py
+++ b/vllm/model_executor/models/starcoder2.py
@@ -91,7 +91,6 @@ class Starcoder2Attention(nn.Module):
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.rope_theta = config.rope_theta
self.max_position_embeddings = config.max_position_embeddings
self.use_bias = config.use_bias
@@ -115,7 +114,7 @@ class Starcoder2Attention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=self.max_position_embeddings,
- base=int(self.rope_theta),
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
self.attn = Attention(
diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py
index 4fff356b29e28..3c377a2c539df 100644
--- a/vllm/model_executor/models/step3_text.py
+++ b/vllm/model_executor/models/step3_text.py
@@ -36,6 +36,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.step3_vl import Step3TextConfig
from .interfaces import SupportsPP
from .utils import (
@@ -144,9 +145,8 @@ class Step3TextAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
norm_eps: float,
- rope_theta: int,
+ rope_parameters: dict[str, Any],
share_q_dim: int | None = None,
- rope_scaling: dict[str, Any] | None = None,
max_position_embedding: int = 8192,
head_dim: int = 256,
cache_config: CacheConfig | None = None,
@@ -198,8 +198,7 @@ class Step3TextAttention(nn.Module):
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position_embedding,
- base=rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=rope_parameters,
)
scaling = self.head_dim**-0.5
self.attn = Attention(
@@ -227,15 +226,13 @@ class Step3TextAttention(nn.Module):
class Step3TextDecoderLayer(nn.Module):
def __init__(
self,
- config: ModelConfig,
+ config: Step3TextConfig,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
- config = config.hf_config
self.hidden_size = config.hidden_size
- rope_scaling = getattr(config, "rope_scaling", None)
self.self_attn = Step3TextAttention(
hidden_size=self.hidden_size,
@@ -247,8 +244,7 @@ class Step3TextDecoderLayer(nn.Module):
max_position_embedding=config.max_position_embedding,
head_dim=config.head_dim,
share_q_dim=config.share_q_dim,
- rope_theta=config.rope_theta,
- rope_scaling=rope_scaling,
+ rope_parameters=config.rope_parameters,
prefix=f"{prefix}.self_attn",
)
@@ -338,7 +334,7 @@ class Step3TextModel(nn.Module):
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
lambda prefix: Step3TextDecoderLayer(
- config=vllm_config.model_config,
+ config=config,
cache_config=cache_config,
quant_config=quant_config,
prefix=prefix,
diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py
index 5d16be1eb3128..1c60cb4148121 100644
--- a/vllm/model_executor/models/step3_vl.py
+++ b/vllm/model_executor/models/step3_vl.py
@@ -20,6 +20,7 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
@@ -667,7 +668,7 @@ class Step3VisionEmbeddings(nn.Module):
self.class_embedding = nn.Parameter(torch.randn(1, self.embed_dim))
- self.patch_embedding = nn.Conv2d(
+ self.patch_embedding = Conv2dLayer(
in_channels=config.num_channels,
out_channels=self.embed_dim,
kernel_size=self.patch_size,
@@ -950,13 +951,13 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
prefix=maybe_prefix(prefix, "vision_model"),
use_data_parallel=self.use_data_parallel,
)
- self.vit_downsampler = nn.Conv2d(
+ self.vit_downsampler = Conv2dLayer(
config.vision_config.hidden_size,
config.vision_config.output_hidden_size,
kernel_size=2,
stride=config.understand_projector_stride,
)
- self.vit_downsampler2 = nn.Conv2d(
+ self.vit_downsampler2 = Conv2dLayer(
config.vision_config.output_hidden_size,
config.vision_config.output_hidden_size * 2,
kernel_size=3,
diff --git a/vllm/model_executor/models/teleflm.py b/vllm/model_executor/models/teleflm.py
index 8a0bec9dff848..bebd7bcaa9249 100644
--- a/vllm/model_executor/models/teleflm.py
+++ b/vllm/model_executor/models/teleflm.py
@@ -74,5 +74,5 @@ class TeleFLMForCausalLM(LlamaForCausalLM):
self.output_mult = self.config.output_mult / self.mup_scale_factor
logit_scale = self.output_mult
self.logits_processor = LogitsProcessor(
- self.unpadded_vocab_size, self.config.vocab_size, logit_scale
+ self.config.vocab_size, scale=logit_scale
)
diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py
index 4973014c3d4ed..31db9d682bd40 100644
--- a/vllm/model_executor/models/transformers/moe.py
+++ b/vllm/model_executor/models/transformers/moe.py
@@ -256,7 +256,14 @@ class MoEMixin(MixtureOfExperts):
def _recursive_replace(module: nn.Module, prefix: str):
for child_name, child_module in module.named_children():
qual_name = maybe_prefix(prefix, child_name)
- if child_name == "experts" and isinstance(child_module, nn.ModuleList):
+ # Naive implementations will have experts as ModuleList
+ is_modulelist = isinstance(child_module, nn.ModuleList)
+ # Packed implementations will have experts as 3D tensors of shapes like:
+ # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
+ # down_proj = (num_experts, intermediate_size, hidden_size)
+ params = list(child_module.parameters())
+ is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
+ if child_name == "experts" and (is_modulelist or is_3d):
# Alias for readability
mlp = module
experts = child_module
diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py
index 517eb54d53ac6..b807f45b5d52b 100644
--- a/vllm/model_executor/models/transformers/utils.py
+++ b/vllm/model_executor/models/transformers/utils.py
@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Literal
import torch
from torch import nn
+from transformers.configuration_utils import ALLOWED_LAYER_TYPES
from vllm.config.utils import getattr_iter
from vllm.logger import init_logger
@@ -203,5 +204,10 @@ def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool:
"""
text_config = vllm_config.model_config.hf_config.get_text_config()
# Dynamic rope scaling is not compatible with torch.compile
- rope_scaling: dict = getattr(text_config, "rope_scaling", None) or {}
- return rope_scaling.get("rope_type") != "dynamic"
+ rope_parameters: dict | None = getattr(text_config, "rope_parameters", None) or {}
+ if rope_parameters:
+ # Nest rope_parameters if not nested already to simplify logic
+ if not set(rope_parameters.keys()).issubset(ALLOWED_LAYER_TYPES):
+ rope_parameters = {"": rope_parameters}
+ return all(rp["rope_type"] != "dynamic" for rp in rope_parameters.values())
+ return True
diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py
index ca5af358e2eed..ccefd7e66697f 100644
--- a/vllm/model_executor/models/utils.py
+++ b/vllm/model_executor/models/utils.py
@@ -21,6 +21,9 @@ from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
)
+from vllm.model_executor.model_loader.online_quantization import (
+ support_quantized_model_reload_from_hp_weights,
+)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.interfaces import supports_any_eagle
from vllm.multimodal import NestedTensors
@@ -316,6 +319,7 @@ class AutoWeightsLoader:
)
raise ValueError(msg)
+ @support_quantized_model_reload_from_hp_weights
def load_weights(
self,
weights: Iterable[tuple[str, torch.Tensor]],
diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py
index 91a10b95a08c0..50587c627160d 100644
--- a/vllm/model_executor/models/whisper.py
+++ b/vllm/model_executor/models/whisper.py
@@ -599,15 +599,16 @@ class WhisperModel(nn.Module):
def forward(
self,
- input_features: torch.Tensor | list[torch.Tensor] | None,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ encoder_outputs: list[torch.Tensor],
) -> torch.Tensor:
- encoder_outputs = self.get_encoder_outputs(input_features)
+ assert len(encoder_outputs) in (0, 1)
+ enc_states = encoder_outputs[0] if len(encoder_outputs) == 1 else None
decoder_outputs = self.decoder(
input_ids=input_ids,
positions=positions,
- encoder_hidden_states=encoder_outputs,
+ encoder_hidden_states=enc_states,
)
return decoder_outputs
@@ -894,13 +895,15 @@ class WhisperForConditionalGeneration(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
+ encoder_outputs: list[torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
- audio_input = self._parse_and_validate_audio_input(**kwargs)
+ if encoder_outputs is None:
+ encoder_outputs = []
decoder_outputs = self.model(
- input_features=audio_input["input_features"],
input_ids=input_ids,
positions=positions,
+ encoder_outputs=encoder_outputs,
)
return decoder_outputs
diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py
index 64e6979c8fcfb..653b5b9beef7b 100644
--- a/vllm/model_executor/models/zamba2.py
+++ b/vllm/model_executor/models/zamba2.py
@@ -128,7 +128,6 @@ class Zamba2Attention(nn.Module):
tp_size = get_tensor_model_parallel_world_size()
self.config = config
self.num_hybrid_layers = num_hybrid_layers
- self.rope_theta = config.rope_theta
self.attention_hidden_size = config.attention_hidden_size
self.total_num_attention_heads = config.num_attention_heads
@@ -233,8 +232,7 @@ class Zamba2Attention(nn.Module):
head_size=self.attention_head_dim,
rotary_dim=self.attention_head_dim,
max_position=config.max_position_embeddings,
- base=self.rope_theta,
- rope_scaling=None,
+ rope_parameters=config.rope_parameters,
is_neox_style=True,
)
@@ -567,11 +565,7 @@ class Zamba2MambaDecoderLayer(nn.Module):
hidden_states = self.input_layernorm(hidden_states)
# Process through Mamba mixer
- output = torch.empty_like(hidden_states)
- self.mamba(
- hidden_states,
- output,
- )
+ output = self.mamba(hidden_states)
# residual connection after mamba
hidden_states = residual + output
diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py
index 759b809433b14..8aad59e84ff25 100644
--- a/vllm/model_executor/utils.py
+++ b/vllm/model_executor/utils.py
@@ -10,7 +10,7 @@ import torch
from vllm.utils.torch_utils import is_torch_equal_or_newer
-def set_random_seed(seed: int) -> None:
+def set_random_seed(seed: int | None) -> None:
from vllm.platforms import current_platform
current_platform.seed_everything(seed)
diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py
index 53052ddc6343c..b93a42ffd24c1 100644
--- a/vllm/multimodal/audio.py
+++ b/vllm/multimodal/audio.py
@@ -7,6 +7,8 @@ from typing import Literal
import numpy as np
import numpy.typing as npt
+import pybase64
+import torch
from vllm.utils.import_utils import PlaceholderModule
@@ -116,3 +118,25 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):
data = buffer.getvalue()
return base64.b64encode(data).decode("utf-8")
+
+
+class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):
+ def __init__(self) -> None:
+ super().__init__()
+
+ def load_bytes(self, data: bytes) -> torch.Tensor:
+ buffer = BytesIO(data)
+ return torch.load(buffer, weights_only=True)
+
+ def load_base64(self, media_type: str, data: str) -> torch.Tensor:
+ return self.load_bytes(pybase64.b64decode(data, validate=True))
+
+ def load_file(self, filepath: Path) -> torch.Tensor:
+ return torch.load(filepath, weights_only=True)
+
+ def encode_base64(self, media: torch.Tensor) -> str:
+ buffer = BytesIO()
+ torch.save(media, buffer)
+ buffer.seek(0)
+ binary_data = buffer.read()
+ return pybase64.b64encode(binary_data).decode("utf-8")
diff --git a/vllm/multimodal/evs.py b/vllm/multimodal/evs.py
index 4a288d2d238c2..8a36ea415da4d 100644
--- a/vllm/multimodal/evs.py
+++ b/vllm/multimodal/evs.py
@@ -185,7 +185,7 @@ def recompute_mrope_positions(
Args:
input_ids: (N,) All input tokens of the prompt (entire sequence).
- multimodal_positions: List of mrope positsions for each media.
+ multimodal_positions: List of mrope positions for each media.
mrope_positions: Existing mrope positions (4, N) for entire sequence.
num_computed_tokens: A number of computed tokens so far.
vision_start_token_id: Token indicating start of vision media.
diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py
index 3f55c46ca334d..1020554e2e073 100644
--- a/vllm/multimodal/utils.py
+++ b/vllm/multimodal/utils.py
@@ -3,7 +3,7 @@
import asyncio
import atexit
-from collections.abc import Iterable, Set
+from collections.abc import Generator, Set
from concurrent.futures import ThreadPoolExecutor
from itertools import groupby
from pathlib import Path
@@ -22,7 +22,7 @@ from vllm.logger import init_logger
from vllm.utils.jsontree import json_map_leaves
from vllm.utils.registry import ExtensionManager
-from .audio import AudioMediaIO
+from .audio import AudioEmbeddingMediaIO, AudioMediaIO
from .base import MediaIO
from .image import ImageEmbeddingMediaIO, ImageMediaIO
from .video import VideoMediaIO
@@ -342,6 +342,17 @@ class MediaConnector:
return image_embedding_io.load_base64("", data)
+ def fetch_audio_embedding(
+ self,
+ data: str,
+ ) -> torch.Tensor:
+ """
+ Load audio embedding from a URL.
+ """
+ audio_embedding_io = AudioEmbeddingMediaIO()
+
+ return audio_embedding_io.load_base64("", data)
+
def encode_audio_base64(
audio: np.ndarray,
@@ -403,7 +414,7 @@ def group_mm_kwargs_by_modality(
pin_memory: bool = False,
merge_by_field_config: bool | None = None,
multimodal_cpu_fields: Set[str] = frozenset(),
-) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
+) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]:
"""Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
modality together into the same `MultiModalKwargs` instance.
diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py
index 369c5e6cb4d10..763f90fde7b6d 100644
--- a/vllm/multimodal/video.py
+++ b/vllm/multimodal/video.py
@@ -63,6 +63,58 @@ class VideoLoader:
) -> tuple[npt.NDArray, dict[str, Any]]:
raise NotImplementedError
+ @staticmethod
+ def _read_frames(
+ cap,
+ frame_indices: set[int],
+ num_expected_frames: int,
+ max_frame_idx: int,
+ ) -> tuple[npt.NDArray, int, list[int]]:
+ import cv2
+
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8)
+
+ i = 0
+ valid_frame_indices = []
+ for idx in range(max_frame_idx + 1):
+ ok = cap.grab()
+ if not ok:
+ # Frame is broken/unreadable, log warning
+ if idx in frame_indices:
+ logger.warning(
+ "Failed to grab frame %d during video loading. "
+ "This frame will be skipped.",
+ idx,
+ )
+ continue
+ if idx in frame_indices:
+ ret, frame = cap.retrieve()
+ if ret:
+ frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+ valid_frame_indices.append(idx)
+ i += 1
+ else:
+ # retrieve() failed even though grab() succeeded
+ logger.warning(
+ "Failed to retrieve frame %d during video loading. "
+ "This frame will be skipped.",
+ idx,
+ )
+
+ valid_num_frames = len(valid_frame_indices)
+ if valid_num_frames < num_expected_frames:
+ logger.warning(
+ "Video loading completed with %d broken/unreadable frames. "
+ "Expected %d frames but only loaded %d frames.",
+ num_expected_frames - valid_num_frames,
+ num_expected_frames,
+ valid_num_frames,
+ )
+
+ return frames[:valid_num_frames], valid_num_frames, valid_frame_indices
+
VIDEO_LOADER_REGISTRY = ExtensionManager()
@@ -120,24 +172,10 @@ class OpenCVVideoBackend(VideoLoader):
)
frame_idx = uniform_sampled_frames.tolist()
- width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
- height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
- frames = np.empty((len(frame_idx), height, width, 3), dtype=np.uint8)
-
- i = 0
- for idx in range(max(frame_idx) + 1):
- ok = cap.grab()
- if not ok:
- break
- if idx in frame_idx:
- ret, frame = cap.retrieve()
- if ret:
- frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
- i += 1
-
- assert i == num_frames_to_sample, (
- f"Expected reading {num_frames_to_sample} frames, "
- f"but only loaded {i} frames from video."
+ # Convert to set for O(1) lookup performance
+ frame_idx_set = set(frame_idx)
+ frames, valid_num_frames, valid_frame_indices = cls._read_frames(
+ cap, frame_idx_set, num_frames_to_sample, max(frame_idx)
)
# Use transformers transformers.video_utils.VideoMetadata format
@@ -148,10 +186,10 @@ class OpenCVVideoBackend(VideoLoader):
"fps": original_fps,
"duration": duration,
"video_backend": "opencv",
- "frames_indices": list(frame_idx),
+ "frames_indices": valid_frame_indices,
# extra field used to control hf processor's video
# sampling behavior
- "do_sample_frames": num_frames_to_sample == total_frames_num,
+ "do_sample_frames": valid_num_frames == total_frames_num,
}
return frames, metadata
@@ -185,10 +223,10 @@ class OpenCVDynamicVideoBackend(OpenCVVideoBackend):
# Refer to:
# https://github.com/huggingface/transformers/blob/v4.55.4/src/transformers/models/glm4v/video_processing_glm4v.py#L103-L140
- frame_indices: range | list[int]
+ frame_indices_list: list[int]
if duration <= max_duration:
n = int(math.floor(duration * fps))
- frame_indices = sorted(
+ frame_indices_list = sorted(
{
min(max_frame_idx, int(math.ceil(i * original_fps / fps)))
for i in range(n)
@@ -197,34 +235,23 @@ class OpenCVDynamicVideoBackend(OpenCVVideoBackend):
else:
num_samples = int(max_duration * fps)
if num_samples >= total_frames_num:
- frame_indices = range(total_frames_num)
+ frame_indices_list = list(range(total_frames_num))
else:
target_seconds = np.linspace(0, duration, num_samples, endpoint=True)
- frame_indices = sorted(
+ frame_indices_list = sorted(
{
min(max_frame_idx, int(math.ceil(t * original_fps)))
for t in target_seconds
}
)
- width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
- height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
- frames = np.empty((len(frame_indices), height, width, 3), dtype=np.uint8)
-
- i = 0
- for idx in range(total_frames_num):
- ok = cap.grab()
- if not ok:
- break
- if idx in frame_indices:
- ret, frame = cap.retrieve()
- if ret:
- frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
- i += 1
-
- assert i == len(frame_indices), (
- f"Expected reading {len(frame_indices)} frames, "
- f"but only loaded {i} frames from video."
+ # Convert to set for O(1) lookup performance
+ frame_indices_set = set(frame_indices_list)
+ frames, valid_num_frames, valid_frame_indices = cls._read_frames(
+ cap,
+ frame_indices_set,
+ len(frame_indices_list),
+ total_frames_num - 1,
)
# Use transformers transformers.video_utils.VideoMetadata format
@@ -233,7 +260,7 @@ class OpenCVDynamicVideoBackend(OpenCVVideoBackend):
"fps": original_fps,
"duration": duration,
"video_backend": "opencv_dynamic",
- "frames_indices": list(frame_indices),
+ "frames_indices": valid_frame_indices,
"do_sample_frames": False,
}
diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py
index 2e4dd8bb808b4..06793a3d1bb14 100644
--- a/vllm/platforms/cuda.py
+++ b/vllm/platforms/cuda.py
@@ -267,25 +267,17 @@ class CudaPlatformBase(Platform):
) -> "AttentionBackendEnum":
from vllm.attention.backends.registry import AttentionBackendEnum
- # For Blackwell GPUs, force TORCH_SDPA for now.
- # See https://github.com/facebookresearch/xformers/issues/1317#issuecomment-3199392579 # noqa: E501
- if cls.has_device_capability(100):
- return AttentionBackendEnum.TORCH_SDPA
-
- if dtype not in (torch.float16, torch.bfloat16):
- return AttentionBackendEnum.XFORMERS
-
- if cls.has_device_capability(80):
+ # Try FlashAttention first
+ try:
backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()
if backend_class.supports_head_size(
head_size
) and backend_class.supports_dtype(dtype):
return AttentionBackendEnum.FLASH_ATTN
- else:
- return AttentionBackendEnum.XFORMERS
- else:
- # Fallback for Volta/Turing GPUs or FA not supported
- return AttentionBackendEnum.XFORMERS
+ except ImportError:
+ pass
+
+ return AttentionBackendEnum.TORCH_SDPA
@classmethod
def get_valid_backends(
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
index 788f9d69c357a..f9005fd7d044c 100644
--- a/vllm/platforms/rocm.py
+++ b/vllm/platforms/rocm.py
@@ -185,6 +185,9 @@ class RocmPlatform(Platform):
"petit_nvfp4",
"torchao",
]
+ # bitsandbytes not supported on gfx9 (warp size 64 limitation)
+ if not on_gfx9():
+ supported_quantization += ["bitsandbytes"]
@classmethod
def get_vit_attn_backend(
@@ -222,7 +225,15 @@ class RocmPlatform(Platform):
from vllm.attention.backends.registry import AttentionBackendEnum
if use_sparse:
- raise NotImplementedError("Sparse Attention is not supported on ROCm.")
+ if kv_cache_dtype.startswith("fp8"):
+ raise ValueError(
+ "ROCMAiterMLASparseBackend doesn't support fp8 kv_cache_dtype."
+ )
+ assert block_size == 1, (
+ "Sparse MLA backend on ROCm only supports block size 1 for now."
+ )
+ logger.info_once("Using Sparse MLA backend on V1 engine.")
+ return AttentionBackendEnum.ROCM_AITER_MLA_SPARSE.get_path()
if use_mla:
if selected_backend is None:
@@ -231,7 +242,6 @@ class RocmPlatform(Platform):
if rocm_aiter_ops.is_mla_enabled() or block_size == 1
else AttentionBackendEnum.TRITON_MLA
)
-
if selected_backend == AttentionBackendEnum.TRITON_MLA:
if block_size != 1:
logger.info_once("Using Triton MLA backend.")
@@ -243,6 +253,9 @@ class RocmPlatform(Platform):
if selected_backend == AttentionBackendEnum.ROCM_AITER_MLA:
logger.info("Using AITER MLA backend.")
return AttentionBackendEnum.ROCM_AITER_MLA.get_path()
+ if selected_backend == AttentionBackendEnum.ROCM_AITER_TRITON_MLA:
+ logger.info("Using AITER TRITON MLA backend.")
+ return AttentionBackendEnum.ROCM_AITER_TRITON_MLA.get_path()
raise ValueError(
f" The selected backend, {selected_backend.name},"
diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py
index 65516827a16da..18a3186b142f1 100644
--- a/vllm/platforms/xpu.py
+++ b/vllm/platforms/xpu.py
@@ -251,10 +251,6 @@ class XPUPlatform(Platform):
) -> None:
"""Copy blocks from src_cache to dst_cache on XPU."""
_src_cache = src_cache[:, src_block_indices]
- if _src_cache.shape[2:] != dst_cache.shape[2:]:
- # To support TP_ratio, HOST KV might be initiated with HND
- # while XPU device KV is with NHD
- _src_cache = _src_cache.permute(0, 1, 3, 2, 4)
dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
@classmethod
@@ -267,8 +263,4 @@ class XPUPlatform(Platform):
) -> None:
"""Copy blocks from XPU to host (CPU)."""
_src_cache = src_cache[:, src_block_indices]
- if _src_cache.shape[2:] != dst_cache.shape[2:]:
- # XPU device KV is with NHD while HOST KV
- # might be initiated with HND for TP_ratio support
- _src_cache = _src_cache.permute(0, 1, 3, 2, 4)
dst_cache[:, dst_block_indices] = _src_cache.cpu()
diff --git a/vllm/plugins/__init__.py b/vllm/plugins/__init__.py
index 0d8988f27959f..4c59d5364a763 100644
--- a/vllm/plugins/__init__.py
+++ b/vllm/plugins/__init__.py
@@ -17,6 +17,9 @@ IO_PROCESSOR_PLUGINS_GROUP = "vllm.io_processor_plugins"
# Platform plugins group will be loaded in all processes when
# `vllm.platforms.current_platform` is called and the value not initialized,
PLATFORM_PLUGINS_GROUP = "vllm.platform_plugins"
+# Stat logger plugins group will be loaded in process0 only when serve vLLM with
+# async mode.
+STAT_LOGGER_PLUGINS_GROUP = "vllm.stat_logger_plugins"
# make sure one process only loads plugins once
plugins_loaded = False
diff --git a/vllm/profiler/gpu_profiler.py b/vllm/profiler/gpu_profiler.py
index 58c6689531615..3e2cbe7296e9d 100644
--- a/vllm/profiler/gpu_profiler.py
+++ b/vllm/profiler/gpu_profiler.py
@@ -1,37 +1,213 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from abc import ABC, abstractmethod
+from contextlib import nullcontext
+
+import torch
+from typing_extensions import override
+
+import vllm.envs as envs
from vllm.logger import init_logger
logger = init_logger(__name__)
-class CudaProfilerWrapper:
+class WorkerProfiler(ABC):
def __init__(self) -> None:
- self._profiler_running = False
+ self._delay_iters = envs.VLLM_PROFILER_DELAY_ITERS
+ if self._delay_iters > 0:
+ logger.info_once(
+ "GPU profiling will start "
+ f"{self._delay_iters} steps after start_profile."
+ )
+
+ self._max_iters = envs.VLLM_PROFILER_MAX_ITERS
+ if self._max_iters > 0:
+ logger.info_once(
+ "GPU profiling will stop "
+ f"after {self._max_iters} worker steps, "
+ "or when stop_profile is received."
+ )
+
+ # Track when the profiler gets triggered by start_profile
+ self._active_iteration_count = 0
+ self._active = False
+
+ # Track when the profiler is actually running
+ self._profiling_for_iters = 0
+ self._running = False
+
+ @abstractmethod
+ def _start(self) -> None:
+ """Start the profiler."""
+ pass
+
+ @abstractmethod
+ def _stop(self) -> None:
+ """Stop the profiler."""
+ pass
+
+ def _call_start(self) -> None:
+ """Call _start with error handling but no safeguards."""
+ try:
+ self._start()
+ self._running = True # Only mark as running if start succeeds
+ except Exception as e:
+ logger.warning("Failed to start profiler: %s", e)
+
+ def _call_stop(self) -> None:
+ """Call _stop with error handling but no safeguards."""
+ try:
+ self._stop()
+ logger.info("Profiler stopped successfully.")
+ except Exception as e:
+ logger.warning("Failed to stop profiler: %s", e)
+ self._running = False # Always mark as not running, assume stop worked
+
+ def start(self) -> None:
+ """Attempt to start the profiler, accounting for delayed starts."""
+ if self._active:
+ logger.debug(
+ "start_profile received when profiler is already active. "
+ "Ignoring request."
+ )
+ return
+ self._active = True
+ if self._delay_iters == 0:
+ self._call_start()
+
+ def step(self) -> None:
+ """Update the profiler state at each worker step,
+ to handle delayed starts and max iteration limits."""
+ if not self._active:
+ return
+
+ self._active_iteration_count += 1
+
+ if (
+ not self._running
+ and self._delay_iters > 0
+ and self._active_iteration_count == self._delay_iters
+ ):
+ logger.info("Starting profiler after delay...")
+ self._call_start()
+
+ if self._running:
+ self._profiling_for_iters += 1
+
+ if (
+ self._max_iters > 0
+ and self._running
+ and self._profiling_for_iters > self._max_iters
+ ):
+ # Automatically stop the profiler after max iters
+ # will be marked as not running, but leave as active so that stop
+ # can clean up properly
+ logger.info("Max profiling iterations reached. Stopping profiler...")
+ self._call_stop()
+ return
+
+ def stop(self) -> None:
+ """Attempt to stop the profiler, accounting for overlapped calls."""
+ if not self._active:
+ logger.debug(
+ "stop_profile received when profiler is not active. Ignoring request."
+ )
+ return
+ self._active = False
+ self._active_iteration_count = 0
+ self._profiling_for_iters = 0
+
+ if self._running:
+ self._call_stop()
+
+ def shutdown(self) -> None:
+ """Ensure profiler is stopped when shutting down."""
+ logger.info_once("Shutting down profiler")
+ if self._running:
+ self.stop()
+
+ def annotate_context_manager(self, name: str):
+ """Return a context manager to annotate profiler traces."""
+ return nullcontext()
+
+
+class TorchProfilerWrapper(WorkerProfiler):
+ def __init__(self, worker_name: str, local_rank: int) -> None:
+ super().__init__()
+
+ self.local_rank = local_rank
+ torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
+ if local_rank in (None, 0):
+ logger.info(
+ "Torch profiling enabled. Traces will be saved to: %s",
+ torch_profiler_trace_dir,
+ )
+ logger.debug(
+ "Profiler config: record_shapes=%s,"
+ "profile_memory=%s,with_stack=%s,with_flops=%s",
+ envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
+ envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
+ envs.VLLM_TORCH_PROFILER_WITH_STACK,
+ envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
+ )
+ self.profiler = torch.profiler.profile(
+ activities=[
+ torch.profiler.ProfilerActivity.CPU,
+ torch.profiler.ProfilerActivity.CUDA,
+ ],
+ record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
+ profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
+ with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
+ with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
+ on_trace_ready=torch.profiler.tensorboard_trace_handler(
+ torch_profiler_trace_dir, worker_name=worker_name, use_gzip=True
+ ),
+ )
+
+ @override
+ def _start(self) -> None:
+ self.profiler.start()
+
+ @override
+ def _stop(self) -> None:
+ self.profiler.stop()
+
+ rank = self.local_rank
+ profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
+ profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
+ sort_key = "self_cuda_time_total"
+ table = self.profiler.key_averages().table(sort_by=sort_key)
+
+ with open(profiler_out_file, "w") as f:
+ print(table, file=f)
+
+ # only print profiler results on rank 0
+ if rank == 0:
+ print(table)
+
+ @override
+ def annotate_context_manager(self, name: str):
+ return torch.profiler.record_function(name)
+
+
+class CudaProfilerWrapper(WorkerProfiler):
+ def __init__(self) -> None:
+ super().__init__()
# Note: lazy import to avoid dependency issues if CUDA is not available.
import torch.cuda.profiler as cuda_profiler
self._cuda_profiler = cuda_profiler
- def start(self) -> None:
- try:
- self._cuda_profiler.start()
- self._profiler_running = True
- logger.info_once("Started CUDA profiler")
- except Exception as e:
- logger.warning_once("Failed to start CUDA profiler: %s", e)
+ @override
+ def _start(self) -> None:
+ self._cuda_profiler.start()
- def stop(self) -> None:
- if self._profiler_running:
- try:
- self._cuda_profiler.stop()
- logger.info_once("Stopped CUDA profiler")
- except Exception as e:
- logger.warning_once("Failed to stop CUDA profiler: %s", e)
- finally:
- self._profiler_running = False
+ @override
+ def _stop(self) -> None:
+ self._cuda_profiler.stop()
- def shutdown(self) -> None:
- """Ensure profiler is stopped when shutting down."""
- self.stop()
+ @override
+ def annotate_context_manager(self, name: str):
+ return torch.cuda.nvtx.range(name)
diff --git a/vllm/ray/lazy_utils.py b/vllm/ray/lazy_utils.py
index 64b5f51571a35..06c91cc3943ae 100644
--- a/vllm/ray/lazy_utils.py
+++ b/vllm/ray/lazy_utils.py
@@ -10,6 +10,8 @@ def is_ray_initialized():
return ray.is_initialized()
except ImportError:
return False
+ except AttributeError:
+ return False
def is_in_ray_actor():
@@ -24,3 +26,5 @@ def is_in_ray_actor():
)
except ImportError:
return False
+ except AttributeError:
+ return False
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
index 901d661634527..fbbe3d4cabb9a 100644
--- a/vllm/sampling_params.py
+++ b/vllm/sampling_params.py
@@ -144,12 +144,6 @@ class SamplingParams(
are generated and streamed cumulatively per request. To see all `n`
outputs upon completion, use `output_kind=RequestOutputKind.FINAL_ONLY`
in `SamplingParams`."""
- best_of: int | None = None
- """Number of output sequences that are generated from the prompt. From
- these `best_of` sequences, the top `n` sequences are returned. `best_of`
- must be greater than or equal to `n`. By default, `best_of` is set to `n`.
- Warning, this is only supported in V0."""
- _real_n: int | None = None
presence_penalty: float = 0.0
"""Penalizes new tokens based on whether they appear in the generated text
so far. Values > 0 encourage the model to use new tokens, while values < 0
@@ -204,6 +198,12 @@ class SamplingParams(
prompt_logprobs: int | None = None
"""Number of log probabilities to return per prompt token.
When set to -1, return all `vocab_size` log probabilities."""
+ flat_logprobs: bool = False
+ """Whether to return logprobs in flatten format (i.e. FlatLogprob)
+ for better performance.
+ NOTE: GC costs of FlatLogprobs is significantly smaller than
+ list[dict[int, Logprob]]. After enabled, PromptLogprobs and
+ SampleLogprobs would populated as FlatLogprobs."""
# NOTE: This parameter is only exposed at the engine level for now.
# It is not exposed in the OpenAI API server, as the OpenAI API does
# not support returning only a list of token IDs.
@@ -259,7 +259,6 @@ class SamplingParams(
@staticmethod
def from_optional(
n: int | None = 1,
- best_of: int | None = None,
presence_penalty: float | None = 0.0,
frequency_penalty: float | None = 0.0,
repetition_penalty: float | None = 1.0,
@@ -309,7 +308,6 @@ class SamplingParams(
return SamplingParams(
n=1 if n is None else n,
- best_of=best_of,
presence_penalty=0.0 if presence_penalty is None else presence_penalty,
frequency_penalty=0.0 if frequency_penalty is None else frequency_penalty,
repetition_penalty=1.0
@@ -342,22 +340,6 @@ class SamplingParams(
)
def __post_init__(self) -> None:
- # how we deal with `best_of`:
- # if `best_of` is not set, we default to `n`;
- # if `best_of` is set, we set `n` to `best_of`,
- # and set `_real_n` to the original `n`.
- # when we return the result, we will check
- # if we need to return `n` or `_real_n` results
- if self.best_of:
- if self.best_of < self.n:
- raise ValueError(
- f"best_of must be greater than or equal to n, "
- f"got n={self.n} and best_of={self.best_of}."
- )
- if not self._real_n:
- self._real_n = self.n
- self.n = self.best_of
-
if 0 < self.temperature < _MAX_TEMP:
logger.warning(
"temperature %s is less than %s, which may cause numerical "
@@ -427,18 +409,6 @@ class SamplingParams(
raise ValueError(f"n must be an int, but is of type {type(self.n)}")
if self.n < 1:
raise ValueError(f"n must be at least 1, got {self.n}.")
- if self.best_of is not None:
- if not isinstance(self.best_of, int):
- raise ValueError(
- f"best_of must be an integer, got {type(self.best_of)}"
- )
- if self.best_of < 1:
- raise ValueError(f"best_of must be at least 1, got {self.best_of}")
- if self.best_of < self.n:
- raise ValueError(
- f"best_of must be greater than or equal to n, "
- f"got n={self.n} and best_of={self.best_of}."
- )
if not -2.0 <= self.presence_penalty <= 2.0:
raise ValueError(
f"presence_penalty must be in [-2, 2], got {self.presence_penalty}."
@@ -513,10 +483,6 @@ class SamplingParams(
"stop strings are only supported when detokenize is True. "
"Set detokenize=True to use stop."
)
- if self.best_of != self._real_n and self.output_kind == (
- RequestOutputKind.DELTA
- ):
- raise ValueError("best_of must equal n to use output_kind=DELTA")
def _verify_greedy_sampling(self) -> None:
if self.n > 1:
diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py
index 49250e071eab2..3d282da8c6112 100644
--- a/vllm/transformers_utils/config.py
+++ b/vllm/transformers_utils/config.py
@@ -1,14 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import fnmatch
import json
import os
import time
from collections.abc import Callable
from dataclasses import asdict
from functools import cache, partial
+from importlib.metadata import version
from pathlib import Path
-from typing import Any, Literal, TypeVar
+from typing import Any, Literal, TypeAlias, TypeVar
import huggingface_hub
from huggingface_hub import (
@@ -24,7 +26,9 @@ from huggingface_hub.utils import (
RepositoryNotFoundError,
RevisionNotFoundError,
)
-from transformers import DeepseekV3Config, GenerationConfig, PretrainedConfig
+from packaging.version import Version
+from transformers import GenerationConfig, PretrainedConfig
+from transformers.configuration_utils import ALLOWED_LAYER_TYPES
from transformers.models.auto.image_processing_auto import get_image_processor_config
from transformers.models.auto.modeling_auto import (
MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
@@ -80,7 +84,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
afmoe="AfmoeConfig",
chatglm="ChatGLMConfig",
deepseek_vl_v2="DeepseekVLV2Config",
- deepseek_v32=DeepseekV3Config,
+ deepseek_v32="DeepseekV3Config",
flex_olmo="FlexOlmoConfig",
kimi_linear="KimiLinearConfig",
kimi_vl="KimiVLConfig",
@@ -200,7 +204,19 @@ class MistralConfigParser(ConfigParserBase):
from vllm.transformers_utils.configs.mistral import adapt_config_dict
- config = adapt_config_dict(config_dict)
+ # Get missing fields from HF config if available
+ try:
+ hf_config_dict, _ = PretrainedConfig.get_config_dict(
+ model,
+ revision=revision,
+ code_revision=code_revision,
+ token=_get_hf_token(),
+ **kwargs,
+ )
+ except OSError: # Not found
+ hf_config_dict = {}
+
+ config = adapt_config_dict(config_dict, defaults=hf_config_dict)
# Mistral configs may define sliding_window as list[int]. Convert it
# to int and add the layer_types list[str] to make it HF compatible
@@ -352,6 +368,41 @@ def list_repo_files(
return with_retry(lookup_files, "Error retrieving file list")
+def list_filtered_repo_files(
+ model_name_or_path: str,
+ allow_patterns: list[str],
+ revision: str | None = None,
+ repo_type: str | None = None,
+ token: str | bool | None = None,
+) -> list[str]:
+ try:
+ all_files = list_repo_files(
+ repo_id=model_name_or_path,
+ revision=revision,
+ token=token,
+ repo_type=repo_type,
+ )
+ except Exception:
+ logger.error(
+ "Error retrieving file list. Please ensure your `model_name_or_path`"
+ "`repo_type`, `token` and `revision` arguments are correctly set. "
+ "Returning an empty list."
+ )
+ return []
+
+ file_list = []
+ # Filter patterns on filenames
+ for pattern in allow_patterns:
+ file_list.extend(
+ [
+ file
+ for file in all_files
+ if fnmatch.fnmatch(os.path.basename(file), pattern)
+ ]
+ )
+ return file_list
+
+
def file_exists(
repo_id: str,
file_name: str,
@@ -390,21 +441,61 @@ def file_or_path_exists(
)
-def patch_rope_scaling(config: PretrainedConfig) -> None:
+def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> None:
+ """Some models may have no rope_theta in their config but still use RoPE.
+ This function sets a default rope_theta if it's missing."""
+ if getattr(config, "rope_parameters", None) is None:
+ config.rope_parameters = {"rope_type": "default"}
+ if "rope_theta" not in config.rope_parameters:
+ config.rope_parameters["rope_theta"] = default_theta
+
+
+def patch_rope_parameters(config: PretrainedConfig) -> None:
"""Provide backwards compatibility for RoPE."""
- text_config = getattr(config, "text_config", None)
- if text_config is not None:
- patch_rope_scaling(text_config)
+ # Retrieve rope_parameters differently based on Transformers version
+ if Version(version("transformers")) >= Version("5.0.0.dev0"):
+ from transformers.modeling_rope_utils import RopeParameters
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is not None:
- patch_rope_scaling_dict(rope_scaling)
+ rope_parameters: RopeParameters | dict[str, RopeParameters] | None = getattr(
+ config, "rope_parameters", None
+ )
+ elif hasattr(config, "rope_parameters"):
+ # We are in Transformers v4 and rope_parameters
+ # has already been patched for this config
+ return
+ else:
+ # Convert Transformers v4 rope_theta and rope_scaling into rope_parameters
+ rope_theta: float | None = getattr(config, "rope_theta", None)
+ rope_scaling: dict | None = getattr(config, "rope_scaling", None)
+ rope_parameters = rope_scaling
+ # Move rope_theta into rope_parameters
+ if rope_theta is not None:
+ rope_parameters = rope_parameters or {"rope_type": "default"}
+ rope_parameters["rope_theta"] = rope_theta
+ # Add original_max_position_embeddings if present
+ if rope_parameters and (
+ ompe := getattr(config, "original_max_position_embeddings", None)
+ ):
+ rope_parameters["original_max_position_embeddings"] = ompe
+ # Write back to config
+ config.rope_parameters = rope_parameters
+
+ # No RoPE parameters to patch
+ if rope_parameters is None:
+ return
+
+ # Handle nested rope_parameters in interleaved sliding attention models
+ if set(rope_parameters.keys()).issubset(ALLOWED_LAYER_TYPES):
+ for rope_parameters_layer_type in rope_parameters.values():
+ patch_rope_parameters_dict(rope_parameters_layer_type)
+ else:
+ patch_rope_parameters_dict(rope_parameters)
-def patch_rope_scaling_dict(rope_scaling: dict[str, Any]) -> None:
- if "rope_type" in rope_scaling and "type" in rope_scaling:
- rope_type = rope_scaling["rope_type"]
- rope_type_legacy = rope_scaling["type"]
+def patch_rope_parameters_dict(rope_parameters: dict[str, Any]) -> None:
+ if "rope_type" in rope_parameters and "type" in rope_parameters:
+ rope_type = rope_parameters["rope_type"]
+ rope_type_legacy = rope_parameters["type"]
if rope_type != rope_type_legacy:
raise ValueError(
f"Found conflicts between 'rope_type={rope_type}' (modern "
@@ -412,28 +503,28 @@ def patch_rope_scaling_dict(rope_scaling: dict[str, Any]) -> None:
"You should only specify one of them."
)
- if "rope_type" not in rope_scaling and "type" in rope_scaling:
- rope_scaling["rope_type"] = rope_scaling["type"]
+ if "rope_type" not in rope_parameters and "type" in rope_parameters:
+ rope_parameters["rope_type"] = rope_parameters["type"]
logger.info("Replacing legacy 'type' key with 'rope_type'")
- if "rope_type" not in rope_scaling:
- raise ValueError("rope_scaling should have a 'rope_type' key")
+ if "rope_type" not in rope_parameters:
+ raise ValueError("rope_parameters should have a 'rope_type' key")
- if rope_scaling["rope_type"] == "su":
- rope_scaling["rope_type"] = "longrope"
+ if rope_parameters["rope_type"] == "su":
+ rope_parameters["rope_type"] = "longrope"
logger.warning("Replacing legacy rope_type 'su' with 'longrope'")
- elif rope_scaling["rope_type"] == "mrope":
- assert "mrope_section" in rope_scaling
- rope_scaling["rope_type"] = "default"
+ elif rope_parameters["rope_type"] == "mrope":
+ assert "mrope_section" in rope_parameters
+ rope_parameters["rope_type"] = "default"
logger.warning("Replacing legacy rope_type 'mrope' with 'default'")
def _uses_mrope(config: PretrainedConfig) -> bool:
- rope_scaling = getattr(config, "rope_scaling", None)
- if rope_scaling is None:
+ rope_parameters = getattr(config, "rope_parameters", None)
+ if rope_parameters is None:
return False
- return "mrope_section" in rope_scaling
+ return "mrope_section" in rope_parameters
def uses_mrope(config: PretrainedConfig) -> bool:
@@ -576,10 +667,14 @@ def get_config(
if config_format == "auto":
try:
- if is_gguf or file_or_path_exists(model, HF_CONFIG_NAME, revision=revision):
- config_format = "hf"
- elif file_or_path_exists(model, MISTRAL_CONFIG_NAME, revision=revision):
+ # First check for Mistral to avoid defaulting to
+ # Transformers implementation.
+ if file_or_path_exists(model, MISTRAL_CONFIG_NAME, revision=revision):
config_format = "mistral"
+ elif is_gguf or file_or_path_exists(
+ model, HF_CONFIG_NAME, revision=revision
+ ):
+ config_format = "hf"
else:
raise ValueError(
"Could not detect config format for no config file found. "
@@ -679,7 +774,14 @@ def get_config(
logger.debug("Overriding HF config with %s", hf_overrides_fn)
config = hf_overrides_fn(config)
- patch_rope_scaling(config)
+ # Exhaustively patch RoPE parameters everywhere they might be
+ patch_rope_parameters(config)
+ patch_rope_parameters(config.get_text_config())
+ SubConfigs: TypeAlias = dict[str, PretrainedConfig]
+ sub_configs: SubConfigs | None = getattr(config, "sub_configs", None)
+ if sub_configs:
+ for sub_config in sub_configs:
+ patch_rope_parameters(getattr(config, sub_config))
if trust_remote_code:
maybe_register_config_serialize_by_value()
diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py
index dcae05a15fec3..d28fd8d033373 100644
--- a/vllm/transformers_utils/configs/__init__.py
+++ b/vllm/transformers_utils/configs/__init__.py
@@ -5,8 +5,13 @@ Model configs may be defined in this directory for the following reasons:
- There is no configuration file defined by HF Hub or Transformers library.
- There is a need to override the existing config to support vLLM.
+- The HF model_type isn't recognized by the Transformers library but can
+ be mapped to an existing Transformers config, such as
+ deepseek-ai/DeepSeek-V3.2-Exp.
"""
+from transformers import DeepseekV3Config
+
from vllm.transformers_utils.configs.afmoe import AfmoeConfig
from vllm.transformers_utils.configs.chatglm import ChatGLMConfig
from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
@@ -44,6 +49,7 @@ __all__ = [
"AfmoeConfig",
"ChatGLMConfig",
"DeepseekVLV2Config",
+ "DeepseekV3Config",
"DotsOCRConfig",
"EAGLEConfig",
"FlexOlmoConfig",
diff --git a/vllm/transformers_utils/configs/afmoe.py b/vllm/transformers_utils/configs/afmoe.py
index 9b634fd037a33..47fee9882f9fc 100644
--- a/vllm/transformers_utils/configs/afmoe.py
+++ b/vllm/transformers_utils/configs/afmoe.py
@@ -24,7 +24,7 @@ class AfmoeConfig(PretrainedConfig):
rms_norm_eps: float = 1e-5,
use_cache: bool = True,
tie_word_embeddings: bool = False,
- rope_theta: float = 10000.0,
+ rope_parameters: dict | None = None,
rope_scaling: dict | None = None,
num_experts: int = 64,
num_experts_per_tok: int = 6,
@@ -56,7 +56,10 @@ class AfmoeConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
+ rope_theta = kwargs.pop("rope_theta", 10000.0)
+ if rope_parameters is None:
+ rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
+ self.rope_parameters = rope_parameters
self.rope_scaling = rope_scaling
self.moe_intermediate_size = moe_intermediate_size
diff --git a/vllm/transformers_utils/configs/arctic.py b/vllm/transformers_utils/configs/arctic.py
index 1707e15285c89..ba4b1a8f701f0 100644
--- a/vllm/transformers_utils/configs/arctic.py
+++ b/vllm/transformers_utils/configs/arctic.py
@@ -85,8 +85,15 @@ class ArcticConfig(PretrainedConfig):
The id of the "end-of-sequence" token.
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
Whether the model's input and output word embeddings should be tied.
- rope_theta (`float`, *optional*, defaults to 1000000.0):
- The base period of the RoPE embeddings.
+ rope_parameters (`dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+ accordingly.
+ Expected contents:
+ `rope_theta` (`float`): The base period of the RoPE embeddings.
+ `rope_type` (`str`):
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+ 'llama3'], with 'default' being the original RoPE implementation.
sliding_window (`int`, *optional*):
Sliding window attention window size. If not specified, will default to `4096`.
attention_dropout (`float`, *optional*, defaults to 0.0):
@@ -132,7 +139,7 @@ class ArcticConfig(PretrainedConfig):
bos_token_id=1,
eos_token_id=2,
tie_word_embeddings=False,
- rope_theta=1e6,
+ rope_parameters: dict[str, Any] | None = None,
sliding_window=None,
attention_dropout=0.0,
num_experts_per_tok=1,
@@ -165,7 +172,10 @@ class ArcticConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
+ rope_theta = kwargs.pop("rope_theta", 1e6)
+ if rope_parameters is None:
+ rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
+ self.rope_parameters = rope_parameters
self.attention_dropout = attention_dropout
self.num_experts_per_tok = num_experts_per_tok
diff --git a/vllm/transformers_utils/configs/flex_olmo.py b/vllm/transformers_utils/configs/flex_olmo.py
index 1f2f4d446288b..c343dc0999a87 100644
--- a/vllm/transformers_utils/configs/flex_olmo.py
+++ b/vllm/transformers_utils/configs/flex_olmo.py
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
from transformers.configuration_utils import PretrainedConfig
@@ -25,8 +26,7 @@ class FlexOlmoConfig(PretrainedConfig):
bos_token_id=None,
eos_token_id=100257,
tie_word_embeddings=False,
- rope_theta=500000.0,
- rope_scaling=None,
+ rope_parameters: dict[str, Any] | None = None,
attention_bias=False,
attention_dropout=0.0,
num_experts_per_tok=5,
@@ -62,8 +62,13 @@ class FlexOlmoConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 500000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
self.num_experts_per_tok = num_experts_per_tok
@@ -73,5 +78,5 @@ class FlexOlmoConfig(PretrainedConfig):
self.norm_topk_prob = norm_topk_prob
# Validate the correctness of rotary position embeddings parameters
# BC: if there is a 'type' field, move it to 'rope_type'.
- if self.rope_scaling is not None and "type" in self.rope_scaling:
- self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+ if self.rope_parameters is not None and "type" in self.rope_parameters:
+ self.rope_parameters["rope_type"] = self.rope_parameters["type"]
diff --git a/vllm/transformers_utils/configs/kimi_linear.py b/vllm/transformers_utils/configs/kimi_linear.py
index 65ddf48c5249b..14894816801d1 100644
--- a/vllm/transformers_utils/configs/kimi_linear.py
+++ b/vllm/transformers_utils/configs/kimi_linear.py
@@ -29,8 +29,7 @@ class KimiLinearConfig(PretrainedConfig):
pad_token_id=0,
bos_token_id=1,
eos_token_id=2,
- rope_theta=10000.0,
- rope_scaling=None,
+ rope_parameters=None,
tie_word_embeddings=False,
moe_intermediate_size: int | None = None,
moe_renormalize: bool = True,
@@ -73,8 +72,13 @@ class KimiLinearConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 10000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
self.q_lora_rank = q_lora_rank
self.kv_lora_rank = kv_lora_rank
diff --git a/vllm/transformers_utils/configs/lfm2_moe.py b/vllm/transformers_utils/configs/lfm2_moe.py
index 37c038e12db80..b399a03c030f0 100644
--- a/vllm/transformers_utils/configs/lfm2_moe.py
+++ b/vllm/transformers_utils/configs/lfm2_moe.py
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
from transformers.configuration_utils import PretrainedConfig
@@ -35,8 +36,8 @@ class Lfm2MoeConfig(PretrainedConfig):
End of stream token id.
tie_word_embeddings (`bool`, *optional*, defaults to `True`):
Whether to tie weight embeddings
- rope_theta (`float`, *optional*, defaults to 1000000.0):
- The base period of the RoPE embeddings.
+ rope_parameters (`dict`, *optional*):
+ The parameters of the RoPE embeddings.
max_position_embeddings (`int`, *optional*, defaults to 128000):
The maximum sequence length that this model might ever be used with.
use_cache (`bool`, *optional*, defaults to `True`):
@@ -100,7 +101,7 @@ class Lfm2MoeConfig(PretrainedConfig):
bos_token_id: int = 1,
eos_token_id: int = 2,
tie_word_embeddings: bool = True,
- rope_theta: float = 1000000.0,
+ rope_parameters: dict[str, Any] | None = None,
max_position_embeddings: int = 128_000,
use_cache: bool = True,
norm_eps: float = 0.00001,
@@ -121,7 +122,10 @@ class Lfm2MoeConfig(PretrainedConfig):
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
- self.rope_theta = rope_theta
+ rope_theta = kwargs.pop("rope_theta", 1000000.0)
+ if rope_parameters is None:
+ rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
+ self.rope_parameters = rope_parameters
self.max_position_embeddings = max_position_embeddings
self.use_cache = use_cache
self.norm_eps = norm_eps
diff --git a/vllm/transformers_utils/configs/midashenglm.py b/vllm/transformers_utils/configs/midashenglm.py
index e49bd26b2b00c..f1bbd057103e4 100644
--- a/vllm/transformers_utils/configs/midashenglm.py
+++ b/vllm/transformers_utils/configs/midashenglm.py
@@ -98,6 +98,6 @@ class MiDashengLMConfig(PretrainedConfig):
if text_config
else Qwen2_5OmniTextConfig()
)
- self.text_config.rope_scaling = None # uses_mrope is false
+ self.text_config.rope_parameters = None # uses_mrope is false
self.audio_token_id = audio_token_id
super().__init__(**kwargs)
diff --git a/vllm/transformers_utils/configs/mistral.py b/vllm/transformers_utils/configs/mistral.py
index c6f04febe37e1..966737aad0867 100644
--- a/vllm/transformers_utils/configs/mistral.py
+++ b/vllm/transformers_utils/configs/mistral.py
@@ -9,14 +9,18 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
-def adapt_config_dict(config_dict: dict[str, Any], **kwargs) -> PretrainedConfig:
- config_dict.update(kwargs)
+def adapt_config_dict(
+ config_dict: dict[str, Any],
+ defaults: dict[str, Any],
+) -> PretrainedConfig:
config_dict = _remap_general_mistral_args(config_dict)
if bool(config_dict.get("quantization")):
config_dict = _remap_mistral_quantization_args(config_dict)
- if bool(config_dict.get("moe")):
+ if config_dict.get("model_type") == "mamba":
+ config_dict["architectures"] = ["Mamba2ForCausalLM"]
+ elif bool(config_dict.get("moe")):
config_dict["architectures"] = ["MixtralForCausalLM"]
else:
config_dict["architectures"] = ["MistralForCausalLM"]
@@ -52,6 +56,9 @@ def adapt_config_dict(config_dict: dict[str, Any], **kwargs) -> PretrainedConfig
if is_audio:
config_dict = _remap_mistral_audio_args(config_dict)
+ for k, v in defaults.items():
+ config_dict.setdefault(k, v)
+
config = PretrainedConfig.from_dict(config_dict)
logger.debug("Initialized config %s", config)
@@ -86,13 +93,17 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
"apply_scale": "apply_yarn_scaling",
}
yarn_config = config.get("yarn") or {}
- config["rope_scaling"] = {
+ config["rope_parameters"] = {
"rope_type": "yarn",
"mscale_all_dim": 1,
}
+
+ if rope_theta := config.pop("rope_theta", None):
+ config["rope_parameters"]["rope_theta"] = rope_theta
+
for old_name, new_name in yarn_config_map.items():
if old_name in yarn_config:
- config["rope_scaling"][new_name] = yarn_config.pop(old_name)
+ config["rope_parameters"][new_name] = yarn_config.pop(old_name)
assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"
@@ -114,7 +125,7 @@ def _remap_general_mistral_args(config: dict) -> dict:
"model_type": ("model_type", "transformer"),
"hidden_act": ("activation", "silu"),
"tie_word_embeddings": ("tied_embeddings", False),
- "max_seq_len": ("max_seq_len", 128_000),
+ "max_seq_len": ("max_seq_len", config.get("max_position_embeddings", 128_000)),
"max_position_embeddings": ("max_position_embeddings", 128_000),
}
diff --git a/vllm/transformers_utils/configs/nemotron.py b/vllm/transformers_utils/configs/nemotron.py
index 60eed549561fb..d112c71d7d20b 100644
--- a/vllm/transformers_utils/configs/nemotron.py
+++ b/vllm/transformers_utils/configs/nemotron.py
@@ -88,8 +88,8 @@ class NemotronConfig(PretrainedConfig):
End of stream token id.
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
Whether to tie weight embeddings
- rope_theta (`float`, *optional*, defaults to 10000.0):
- The base period of the RoPE embeddings.
+ rope_parameters (`dict`, *optional*):
+ The parameters of the RoPE embeddings.
partial_rotary_factor (`float`, *optional*, defaults to 0.5):
Percentage of the query and keys which will have rotary embedding.
attention_bias (`bool`, *optional*, defaults to `False`):
@@ -132,8 +132,7 @@ class NemotronConfig(PretrainedConfig):
bos_token_id=2,
eos_token_id=3,
tie_word_embeddings=False,
- rope_theta=10000.0,
- rope_scaling=None,
+ rope_parameters=None,
partial_rotary_factor=0.5,
attention_bias=False,
attention_dropout=0.0,
@@ -160,8 +159,13 @@ class NemotronConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.norm_eps = norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 10000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
# for backward compatibility
partial_rotary_factor = (
kwargs.get("rope_percent")
@@ -169,7 +173,7 @@ class NemotronConfig(PretrainedConfig):
or partial_rotary_factor
)
self.partial_rotary_factor = partial_rotary_factor
- self._rope_scaling_validation()
+ self._rope_parameters_validation()
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
self.mlp_bias = mlp_bias
@@ -182,31 +186,29 @@ class NemotronConfig(PretrainedConfig):
**kwargs,
)
- def _rope_scaling_validation(self):
+ def _rope_parameters_validation(self):
"""
- Validate the `rope_scaling` configuration.
+ Validate the `rope_parameters` configuration.
"""
- if self.rope_scaling is None:
+ if self.rope_parameters is None:
return
- if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
+ rope_type: str | None = self.rope_parameters.get("rope_type", None)
+ factor: float | None = self.rope_parameters.get("factor", None)
+
+ if rope_type not in {"default", "linear", "dynamic"}:
raise ValueError(
- "`rope_scaling` must be a dictionary with two fields, "
- f"`type` and `factor`, got {self.rope_scaling}"
- )
- rope_scaling_type = self.rope_scaling.get("type", None)
- rope_scaling_factor = self.rope_scaling.get("factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
- raise ValueError(
- "`rope_scaling`'s type field must be one of ['linear', "
- f"'dynamic'], got {rope_scaling_type}"
- )
- if (
- rope_scaling_factor is None
- or not isinstance(rope_scaling_factor, float)
- or rope_scaling_factor <= 1.0
- ):
- raise ValueError(
- "`rope_scaling`'s factor field must be a float > 1, got "
- f"{rope_scaling_factor}"
+ "`rope_type` must be one of ['default', 'linear', 'dynamic'], "
+ f"got {rope_type}"
)
+ if rope_type != "default":
+ if factor is None:
+ raise ValueError(
+ "If `rope_type` is not 'default', `rope_parameters` "
+ "must include a `factor` field. Got `None`."
+ )
+ if not isinstance(factor, float) or factor <= 1.0:
+ raise ValueError(
+ "`rope_parameters`'s factor field must be a float > 1, got "
+ f"{factor}"
+ )
diff --git a/vllm/transformers_utils/configs/olmo3.py b/vllm/transformers_utils/configs/olmo3.py
index f5a9a7cd36bdb..c4691b661af39 100644
--- a/vllm/transformers_utils/configs/olmo3.py
+++ b/vllm/transformers_utils/configs/olmo3.py
@@ -24,8 +24,7 @@ class Olmo3Config(PretrainedConfig):
bos_token_id=None,
eos_token_id=50279,
tie_word_embeddings=False,
- rope_theta=10000.0,
- rope_scaling=None,
+ rope_parameters=None,
attention_bias=False,
attention_dropout=0.0,
rms_norm_eps=1e-5,
@@ -63,8 +62,13 @@ class Olmo3Config(PretrainedConfig):
self.hidden_act = hidden_act
self.initializer_range = initializer_range
self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 10000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
diff --git a/vllm/transformers_utils/configs/qwen3_next.py b/vllm/transformers_utils/configs/qwen3_next.py
index 21750bde2f878..d2fe58d48da6f 100644
--- a/vllm/transformers_utils/configs/qwen3_next.py
+++ b/vllm/transformers_utils/configs/qwen3_next.py
@@ -66,13 +66,12 @@ class Qwen3NextConfig(PretrainedConfig):
relevant if `config.is_decoder=True`.
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
Whether the model's input and output word embeddings should be tied.
- rope_theta (`float`, *optional*, defaults to 10000.0):
- The base period of the RoPE embeddings.
- rope_scaling (`Dict`, *optional*):
+ rope_parameters (`dict`, *optional*):
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
accordingly.
Expected contents:
+ `rope_theta` (`float`): The base period of the RoPE embeddings.
`rope_type` (`str`):
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
'llama3'], with 'default' being the original RoPE implementation.
@@ -199,8 +198,7 @@ class Qwen3NextConfig(PretrainedConfig):
rms_norm_eps=1e-6,
use_cache=True,
tie_word_embeddings=False,
- rope_theta=10000.0,
- rope_scaling=None,
+ rope_parameters=None,
partial_rotary_factor=0.25,
attention_bias=False,
attention_dropout=0.0,
@@ -236,8 +234,13 @@ class Qwen3NextConfig(PretrainedConfig):
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 10000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
self.partial_rotary_factor = partial_rotary_factor
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
diff --git a/vllm/transformers_utils/configs/step3_vl.py b/vllm/transformers_utils/configs/step3_vl.py
index 637b82d88e265..0ee650a70451f 100644
--- a/vllm/transformers_utils/configs/step3_vl.py
+++ b/vllm/transformers_utils/configs/step3_vl.py
@@ -52,8 +52,7 @@ class Step3TextConfig(PretrainedConfig):
moe_intermediate_size: int = 5120,
moe_num_experts: int = 48,
moe_top_k: int = 3,
- rope_theta: float = 500000,
- rope_scaling: dict[str, Any] | None = None,
+ rope_parameters: dict[str, Any] | None = None,
max_position_embedding: int = 65536,
share_expert_dim: int = 5120,
share_q_dim: int = 2048,
@@ -130,8 +129,13 @@ class Step3TextConfig(PretrainedConfig):
self.moe_intermediate_size = moe_intermediate_size
self.moe_num_experts = moe_num_experts
self.moe_top_k = moe_top_k
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
+ rope_theta = kwargs.pop("rope_theta", 500000.0)
+ if "rope_theta" not in rope_parameters:
+ rope_parameters["rope_theta"] = rope_theta
+ self.rope_parameters = rope_parameters
self.max_position_embedding = max_position_embedding
self.share_expert_dim = share_expert_dim
self.share_q_dim = share_q_dim
diff --git a/vllm/transformers_utils/gguf_utils.py b/vllm/transformers_utils/gguf_utils.py
new file mode 100644
index 0000000000000..2bf59c91a3bb1
--- /dev/null
+++ b/vllm/transformers_utils/gguf_utils.py
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GGUF utility functions."""
+
+from pathlib import Path
+
+import gguf
+from gguf.constants import Keys, VisionProjectorType
+from transformers import Gemma3Config, PretrainedConfig, SiglipVisionConfig
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def detect_gguf_multimodal(model: str) -> Path | None:
+ """Check if GGUF model has multimodal projector file.
+
+ Args:
+ model: Model path string
+
+ Returns:
+ Path to mmproj file if found, None otherwise
+ """
+ if not model.endswith(".gguf"):
+ return None
+
+ try:
+ model_path = Path(model)
+ if not model_path.is_file():
+ return None
+
+ model_dir = model_path.parent
+ mmproj_patterns = ["mmproj.gguf", "mmproj-*.gguf", "*mmproj*.gguf"]
+ for pattern in mmproj_patterns:
+ mmproj_files = list(model_dir.glob(pattern))
+ if mmproj_files:
+ return mmproj_files[0]
+ return None
+ except Exception:
+ return None
+
+
+def extract_vision_config_from_gguf(mmproj_path: str) -> "SiglipVisionConfig | None":
+ """Extract vision config parameters from mmproj.gguf metadata.
+
+ Reads vision encoder configuration from GGUF metadata fields using
+ standardized GGUF constants. Automatically detects the projector type
+ (e.g., gemma3, llama4) and applies model-specific parameters accordingly.
+
+ The function extracts standard CLIP vision parameters from GGUF metadata
+ and applies projector-type-specific customizations. For unknown projector
+ types, it uses safe defaults from SiglipVisionConfig.
+
+ Args:
+ mmproj_path: Path to mmproj.gguf file (str or Path)
+
+ Returns:
+ SiglipVisionConfig if extraction succeeds, None if any required
+ field is missing from the GGUF metadata
+
+ Raises:
+ Exception: Exceptions from GGUF reading (file not found, corrupted
+ file, etc.) propagate directly from gguf.GGUFReader
+ """
+ reader = gguf.GGUFReader(str(mmproj_path))
+
+ # Detect projector type to apply model-specific parameters
+ projector_type = None
+ projector_type_field = reader.get_field(Keys.Clip.PROJECTOR_TYPE)
+ if projector_type_field:
+ try:
+ projector_type = bytes(projector_type_field.parts[-1]).decode("utf-8")
+ except (AttributeError, UnicodeDecodeError) as e:
+ logger.warning("Failed to decode projector type from GGUF: %s", e)
+
+ # Map GGUF field constants to SiglipVisionConfig parameters.
+ # Uses official GGUF constants from gguf-py for standardization.
+ # Format: {gguf_constant: (param_name, dtype)}
+ VISION_CONFIG_FIELDS = {
+ Keys.ClipVision.EMBEDDING_LENGTH: ("hidden_size", int),
+ Keys.ClipVision.FEED_FORWARD_LENGTH: ("intermediate_size", int),
+ Keys.ClipVision.BLOCK_COUNT: ("num_hidden_layers", int),
+ Keys.ClipVision.Attention.HEAD_COUNT: ("num_attention_heads", int),
+ Keys.ClipVision.IMAGE_SIZE: ("image_size", int),
+ Keys.ClipVision.PATCH_SIZE: ("patch_size", int),
+ Keys.ClipVision.Attention.LAYERNORM_EPS: ("layer_norm_eps", float),
+ }
+
+ # Extract and validate all required fields
+ config_params = {}
+ for gguf_key, (param_name, dtype) in VISION_CONFIG_FIELDS.items():
+ field = reader.get_field(gguf_key)
+ if field is None:
+ logger.warning(
+ "Missing required vision config field '%s' in mmproj.gguf",
+ gguf_key,
+ )
+ return None
+ # Extract scalar value from GGUF field and convert to target type
+ config_params[param_name] = dtype(field.parts[-1])
+
+ # Apply model-specific parameters based on projector type
+ if projector_type == VisionProjectorType.GEMMA3:
+ # Gemma3 doesn't use the vision pooling head (multihead attention)
+ # This is a vLLM-specific parameter used in SiglipVisionTransformer
+ config_params["vision_use_head"] = False
+ logger.info("Detected Gemma3 projector, disabling vision pooling head")
+ # Add other projector-type-specific customizations here as needed
+ # elif projector_type == VisionProjectorType.LLAMA4:
+ # config_params["vision_use_head"] = ...
+
+ # Create config with extracted parameters
+ # Note: num_channels and attention_dropout use SiglipVisionConfig defaults
+ # (3 and 0.0 respectively) which are correct for all models
+ config = SiglipVisionConfig(**config_params)
+
+ if projector_type:
+ logger.info(
+ "Extracted vision config from mmproj.gguf (projector_type: %s)",
+ projector_type,
+ )
+ else:
+ logger.info("Extracted vision config from mmproj.gguf metadata")
+
+ return config
+
+
+def maybe_patch_hf_config_from_gguf(
+ model: str,
+ hf_config: PretrainedConfig,
+) -> PretrainedConfig:
+ """Patch HF config for GGUF models.
+
+ Applies GGUF-specific patches to HuggingFace config:
+ 1. For multimodal models: patches architecture and vision config
+ 2. For all GGUF models: overrides vocab_size from embedding tensor
+
+ This ensures compatibility with GGUF models that have extended
+ vocabularies (e.g., Unsloth) where the GGUF file contains more
+ tokens than the HuggingFace tokenizer config specifies.
+
+ Args:
+ model: Model path string
+ hf_config: HuggingFace config to patch in-place
+
+ Returns:
+ Updated HuggingFace config
+ """
+ # Patch multimodal config if mmproj.gguf exists
+ mmproj_path = detect_gguf_multimodal(model)
+ if mmproj_path is not None:
+ vision_config = extract_vision_config_from_gguf(str(mmproj_path))
+
+ # Create HF config for Gemma3 multimodal
+ text_config = hf_config.get_text_config()
+ is_gemma3 = hf_config.model_type in ("gemma3", "gemma3_text")
+ if vision_config is not None and is_gemma3:
+ new_hf_config = Gemma3Config.from_text_vision_configs(
+ text_config=text_config,
+ vision_config=vision_config,
+ architectures=["Gemma3ForConditionalGeneration"],
+ )
+ hf_config = new_hf_config
+
+ return hf_config
diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py
index b3469c1b18f2d..8deacb5b07913 100644
--- a/vllm/transformers_utils/processor.py
+++ b/vllm/transformers_utils/processor.py
@@ -18,7 +18,7 @@ from transformers.processing_utils import ProcessorMixin
from transformers.video_processing_utils import BaseVideoProcessor
from typing_extensions import TypeVar
-from vllm.transformers_utils.utils import convert_model_repo_to_path
+from vllm.transformers_utils.utils import check_gguf_file, convert_model_repo_to_path
from vllm.utils.func_utils import get_allowed_kwarg_only_overrides
if TYPE_CHECKING:
@@ -236,9 +236,20 @@ def cached_processor_from_config(
processor_cls: type[_P] | tuple[type[_P], ...] = ProcessorMixin,
**kwargs: Any,
) -> _P:
+ if check_gguf_file(model_config.model):
+ assert not check_gguf_file(model_config.tokenizer), (
+ "For multimodal GGUF models, the original tokenizer "
+ "should be used to correctly load processor."
+ )
+ model = model_config.tokenizer
+ revision = model_config.tokenizer_revision
+ else:
+ model = model_config.model
+ revision = model_config.revision
+
return cached_get_processor_without_dynamic_kwargs(
- model_config.model,
- revision=model_config.revision,
+ model,
+ revision=revision,
trust_remote_code=model_config.trust_remote_code,
processor_cls=processor_cls, # type: ignore[arg-type]
**_merge_mm_kwargs(model_config, processor_cls, **kwargs),
@@ -339,9 +350,19 @@ def cached_image_processor_from_config(
model_config: "ModelConfig",
**kwargs: Any,
):
+ if check_gguf_file(model_config.model):
+ assert not check_gguf_file(model_config.tokenizer), (
+ "For multimodal GGUF models, the original tokenizer "
+ "should be used to correctly load image processor."
+ )
+ model = model_config.tokenizer
+ revision = model_config.tokenizer_revision
+ else:
+ model = model_config.model
+ revision = model_config.revision
return cached_get_image_processor(
- model_config.model,
- revision=model_config.revision,
+ model,
+ revision=revision,
trust_remote_code=model_config.trust_remote_code,
**_merge_mm_kwargs(model_config, AutoImageProcessor, **kwargs),
)
diff --git a/vllm/transformers_utils/tokenizer.py b/vllm/transformers_utils/tokenizer.py
index a393568909d27..233076741503d 100644
--- a/vllm/transformers_utils/tokenizer.py
+++ b/vllm/transformers_utils/tokenizer.py
@@ -3,8 +3,8 @@
import contextlib
import copy
+import importlib.util
import os
-import warnings
from functools import lru_cache
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypeAlias
@@ -15,7 +15,10 @@ from typing_extensions import assert_never
from vllm import envs
from vllm.logger import init_logger
-from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config
+from vllm.transformers_utils.config import (
+ get_sentence_transformer_tokenizer_config,
+ list_filtered_repo_files,
+)
from vllm.transformers_utils.tokenizers import MistralTokenizer
from vllm.transformers_utils.utils import check_gguf_file
@@ -182,25 +185,29 @@ def get_tokenizer(
kwargs["gguf_file"] = Path(tokenizer_name).name
tokenizer_name = Path(tokenizer_name).parent
- # if tokenizer is from official mistral org
- is_from_mistral_org = str(tokenizer_name).split("/")[0] == "mistralai"
- if is_from_mistral_org and tokenizer_mode != "mistral":
- warnings.warn(
- "It is strongly recommended to run mistral models with "
- '`--tokenizer-mode "mistral"` to ensure correct '
- "encoding and decoding.",
- FutureWarning,
- stacklevel=2,
+ # if `tokenizer_mode` == "auto", check if tokenizer can be loaded via Mistral format
+ # first to use official Mistral tokenizer if possible.
+ mistral_common_installed = importlib.util.find_spec("mistral_common") is not None
+ if tokenizer_mode == "auto" and mistral_common_installed:
+ allow_patterns = ["tekken.json", "tokenizer.model.v*"]
+ files_list = list_filtered_repo_files(
+ model_name_or_path=str(tokenizer_name),
+ allow_patterns=allow_patterns,
+ revision=revision,
)
+ if len(files_list) > 0:
+ tokenizer_mode = "mistral"
tokenizer: AnyTokenizer
if tokenizer_mode == "mistral":
+ logger.debug_once(f"Loading MistralTokenizer from {tokenizer_name}")
tokenizer = MistralTokenizer.from_pretrained(
str(tokenizer_name), revision=revision
)
elif tokenizer_mode == "custom":
from vllm.transformers_utils.tokenizer_base import TokenizerRegistry
+ logger.debug_once(f"Loading CustomTokenizer from {tokenizer_name}")
tokenizer = TokenizerRegistry.get_tokenizer(
str(tokenizer_name),
*args,
@@ -210,6 +217,7 @@ def get_tokenizer(
)
else:
try:
+ logger.debug_once(f"Loading AutoTokenizer from {tokenizer_name}")
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_name,
*args,
diff --git a/vllm/transformers_utils/utils.py b/vllm/transformers_utils/utils.py
index 1ae42ba622dc4..901a64d9d2633 100644
--- a/vllm/transformers_utils/utils.py
+++ b/vllm/transformers_utils/utils.py
@@ -27,6 +27,7 @@ def is_cloud_storage(model_or_path: str) -> bool:
return is_s3(model_or_path) or is_gcs(model_or_path)
+@cache
def check_gguf_file(model: str | PathLike) -> bool:
"""Check if the file is a GGUF model."""
model = Path(model)
diff --git a/vllm/utils/__init__.py b/vllm/utils/__init__.py
index 3ef44e7703204..d94da71b289f3 100644
--- a/vllm/utils/__init__.py
+++ b/vllm/utils/__init__.py
@@ -49,7 +49,6 @@ STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"
# Possible string values of STR_BACKEND_ENV_VAR
# register, corresponding to possible backends
STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"
-STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
STR_INVALID_VAL: str = "INVALID"
diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py
index b5ab37534dd78..b25c1e3e1ece3 100644
--- a/vllm/utils/deep_gemm.py
+++ b/vllm/utils/deep_gemm.py
@@ -325,6 +325,7 @@ DEFAULT_BLOCK_SIZE = [128, 128]
def per_block_cast_to_fp8(
x: torch.Tensor, block_size: list[int] = DEFAULT_BLOCK_SIZE, use_ue8m0: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
+ fp8_dtype = current_platform.fp8_dtype()
assert x.dim() == 2
m, n = x.shape
block_m, block_n = block_size
@@ -334,9 +335,9 @@ def per_block_cast_to_fp8(
x_padded[:m, :n] = x
x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
- sf = x_amax / 448.0
+ sf = x_amax / 224.0 if current_platform.is_fp8_fnuz() else x_amax / 448.0
sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
- x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
+ x_scaled = (x_view * (1.0 / sf)).to(fp8_dtype)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
x_view.size(0), x_view.size(2)
)
@@ -365,11 +366,18 @@ def should_use_deepgemm_for_fp8_linear(
):
if supports_deep_gemm is None:
supports_deep_gemm = is_deep_gemm_supported()
+
+ # Verify DeepGEMM N/K dims requirements
+ # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul
+ # test inside kernels/quatization/test_block_fp8.py
+ N_MULTIPLE = 64
+ K_MULTIPLE = 128
+
return (
supports_deep_gemm
and output_dtype == torch.bfloat16
- and weight.shape[0] % 128 == 0
- and weight.shape[1] % 128 == 0
+ and weight.shape[0] % N_MULTIPLE == 0
+ and weight.shape[1] % K_MULTIPLE == 0
)
diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py
index 1209d64901bf5..9f9976d52b4ae 100644
--- a/vllm/utils/flashinfer.py
+++ b/vllm/utils/flashinfer.py
@@ -114,7 +114,17 @@ flashinfer_trtllm_fp8_per_tensor_scale_moe = _lazy_import_wrapper(
flashinfer_cutlass_fused_moe = _lazy_import_wrapper(
"flashinfer.fused_moe", "cutlass_fused_moe"
)
+flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper(
+ "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"
+)
flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
+nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize")
+silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper(
+ "flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize"
+)
+scaled_fp4_grouped_quantize = _lazy_import_wrapper(
+ "flashinfer", "scaled_fp4_grouped_quantize"
+)
nvfp4_block_scale_interleave = _lazy_import_wrapper(
"flashinfer", "nvfp4_block_scale_interleave"
)
@@ -166,6 +176,14 @@ def has_flashinfer_moe() -> bool:
)
+@functools.cache
+def has_flashinfer_cutedsl() -> bool:
+ """Return ``True`` if FlashInfer cutedsl module is available."""
+ return (
+ has_flashinfer() and importlib.util.find_spec("flashinfer.cute_dsl") is not None
+ )
+
+
@functools.cache
def has_flashinfer_cutlass_fused_moe() -> bool:
"""Return `True` if FlashInfer CUTLASS fused MoE is available."""
@@ -187,6 +205,26 @@ def has_flashinfer_cutlass_fused_moe() -> bool:
return True
+@functools.cache
+def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool:
+ """Return ``True`` if FlashInfer CUTLASS fused MoE is available."""
+ if not has_flashinfer_cutedsl():
+ return False
+
+ # Check if all required functions are available
+ required_functions = [
+ ("flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"),
+ ("flashinfer", "scaled_fp4_grouped_quantize"),
+ ("flashinfer", "silu_and_scaled_nvfp4_experts_quantize"),
+ ]
+
+ for module_name, attr_name in required_functions:
+ mod = _get_submodule(module_name)
+ if not mod or not hasattr(mod, attr_name):
+ return False
+ return True
+
+
@functools.cache
def has_nvidia_artifactory() -> bool:
"""Return `True` if NVIDIA's artifactory is accessible.
@@ -472,7 +510,10 @@ __all__ = [
"has_flashinfer",
"flashinfer_trtllm_fp8_block_scale_moe",
"flashinfer_cutlass_fused_moe",
+ "flashinfer_cutedsl_grouped_gemm_nt_masked",
"flashinfer_fp4_quantize",
+ "silu_and_mul_scaled_nvfp4_experts_quantize",
+ "scaled_fp4_grouped_quantize",
"nvfp4_block_scale_interleave",
"trtllm_fp4_block_scale_moe",
"autotune",
@@ -480,6 +521,7 @@ __all__ = [
"has_flashinfer_comm",
"has_flashinfer_all2all",
"has_flashinfer_cutlass_fused_moe",
+ "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
"has_nvidia_artifactory",
"supports_trtllm_attention",
"can_use_trtllm_attention",
diff --git a/vllm/utils/gc_utils.py b/vllm/utils/gc_utils.py
index 160ac9ac263a9..c56b1794230e9 100644
--- a/vllm/utils/gc_utils.py
+++ b/vllm/utils/gc_utils.py
@@ -53,6 +53,7 @@ class GCDebugger:
self.config = config
# Start time in micro second of this GC cycle
self.start_time_ns: int = time.monotonic_ns()
+ self.num_objects: int = 0
# If config.top_objects is positive,
# compute top collected objects by object types
self.gc_top_collected_objects: str = ""
@@ -68,8 +69,10 @@ class GCDebugger:
# Before GC started, record GC start time
# and top collected objects
self.start_time_ns = time.monotonic_ns()
+ objects = gc.get_objects(generation)
+ self.num_objects = len(objects)
self.gc_top_collected_objects = _compute_top_gc_collected_objects(
- gc.get_objects(generation), self.config.top_objects
+ objects, self.config.top_objects
)
elif phase == "stop":
# After GC finished, Record GC elapsed time and
@@ -77,9 +80,10 @@ class GCDebugger:
elpased_ms = (time.monotonic_ns() - self.start_time_ns) / 1e6
logger.info(
"GC took %.3fms to complete. "
- "Collected %s objects in GC generation %d.%s",
+ "Collected %s objects (out of %d) in GC generation %d.%s",
elpased_ms,
str(info.get("collected", "?")),
+ self.num_objects,
generation,
(
f" Top collected objects: \n{self.gc_top_collected_objects}"
diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py
index f01d2c7a6a33d..ff0f0350fd941 100644
--- a/vllm/utils/import_utils.py
+++ b/vllm/utils/import_utils.py
@@ -18,6 +18,10 @@ from typing import Any
import regex as re
from typing_extensions import Never
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
# TODO: This function can be removed if transformer_modules classes are
# serialized by value when communicating between processes
@@ -62,6 +66,35 @@ def import_pynvml():
return pynvml
+@cache
+def import_triton_kernels():
+ """
+ For convenience, prioritize triton_kernels that is available in
+ `site-packages`. Use `vllm.third_party.triton_kernels` as a fall-back.
+ """
+ if _has_module("triton_kernels"):
+ import triton_kernels
+
+ logger.debug_once(
+ f"Loading module triton_kernels from {triton_kernels.__file__}.",
+ scope="local",
+ )
+ elif _has_module("vllm.third_party.triton_kernels"):
+ import vllm.third_party.triton_kernels as triton_kernels
+
+ logger.debug_once(
+ f"Loading module triton_kernels from {triton_kernels.__file__}.",
+ scope="local",
+ )
+ sys.modules["triton_kernels"] = triton_kernels
+ else:
+ logger.info_once(
+ "triton_kernels unavailable in this build. "
+ "Please consider installing triton_kernels from "
+ "https://github.com/triton-lang/triton/tree/main/python/triton_kernels"
+ )
+
+
def import_from_path(module_name: str, file_path: str | os.PathLike):
"""
Import a Python file according to its file path.
@@ -397,7 +430,12 @@ def has_deep_gemm() -> bool:
def has_triton_kernels() -> bool:
"""Whether the optional `triton_kernels` package is available."""
- return _has_module("triton_kernels")
+ is_available = _has_module("triton_kernels") or _has_module(
+ "vllm.third_party.triton_kernels"
+ )
+ if is_available:
+ import_triton_kernels()
+ return is_available
def has_tilelang() -> bool:
diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py
index 5968884e232a4..cc872040b6c5f 100644
--- a/vllm/utils/system_utils.py
+++ b/vllm/utils/system_utils.py
@@ -22,7 +22,7 @@ from .platform_utils import cuda_is_initialized, xpu_is_initialized
logger = init_logger(__name__)
-CYAN = "\033[1;36m"
+CYAN = "\033[0;36m"
RESET = "\033[0;0m"
@@ -142,7 +142,10 @@ def set_process_title(
def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
"""Add colored prefix to file output for log decoration."""
- prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
+ if envs.NO_COLOR:
+ prefix = f"({worker_name} pid={pid}) "
+ else:
+ prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
file_write = file.write
def write_with_prefix(s: str):
diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py
index 7c094e14cff72..3661dfd09047a 100644
--- a/vllm/utils/torch_utils.py
+++ b/vllm/utils/torch_utils.py
@@ -426,8 +426,7 @@ def aux_stream() -> torch.cuda.Stream | None:
from vllm.platforms import current_platform
- # TODO: validate this works properly on ROCm platform.
- if _aux_stream is None and current_platform.is_cuda():
+ if _aux_stream is None and current_platform.is_cuda_alike():
_aux_stream = torch.cuda.Stream()
return _aux_stream
diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py
index f1254352c0585..590bf91b0d057 100644
--- a/vllm/v1/attention/backends/cpu_attn.py
+++ b/vllm/v1/attention/backends/cpu_attn.py
@@ -25,7 +25,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec
logger = init_logger(__name__)
-_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86,)
+_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86, CpuArchEnum.ARM)
class CPUAttentionBackend(AttentionBackend):
@@ -491,6 +491,9 @@ def _get_attn_isa(dtype: torch.dtype, block_size: int) -> str:
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
return "amx"
elif block_size % 32 == 0:
- return "vec"
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+ return "neon"
+ else:
+ return "vec"
else:
return "vec16"
diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py
index a5d4435000d4d..a9a4af5ac1183 100755
--- a/vllm/v1/attention/backends/flash_attn.py
+++ b/vllm/v1/attention/backends/flash_attn.py
@@ -32,7 +32,7 @@ if is_flash_attn_varlen_func_available():
get_scheduler_metadata,
reshape_and_cache_flash,
)
-from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config
from vllm.config.cache import CacheDType
from vllm.distributed.parallel_state import get_dcp_group
from vllm.logger import init_logger
@@ -56,11 +56,26 @@ logger = init_logger(__name__)
class FlashAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- # 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
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [16, 32, 64]
+
+ @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)]
@staticmethod
def get_name() -> str:
@@ -99,12 +114,20 @@ class FlashAttentionBackend(AttentionBackend):
return (2, num_blocks, block_size, num_kv_heads, head_size)
@staticmethod
- def get_kv_cache_stride_order() -> tuple[int, ...]:
+ def get_kv_cache_stride_order(
+ include_num_layers_dimension: bool = False,
+ ) -> tuple[int, ...]:
# `stride_order` indicates the permutation that gets
# us from `get_kv_cache_shape` to the actual memory layout we want.
cache_layout = get_kv_cache_layout()
- if cache_layout == "NHD":
+ if cache_layout == "NHD" and include_num_layers_dimension:
+ # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
+ return (2, 0, 1, 3, 4, 5)
+ elif cache_layout == "NHD":
stride_order = (0, 1, 2, 3, 4)
+ elif cache_layout == "HND" and include_num_layers_dimension:
+ # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size)
+ return (2, 4, 0, 1, 3, 5)
elif cache_layout == "HND":
stride_order = (0, 1, 3, 2, 4)
else:
@@ -119,8 +142,8 @@ class FlashAttentionBackend(AttentionBackend):
raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
@classmethod
- def get_supported_head_sizes(cls) -> list[int]:
- return [32, 64, 96, 128, 160, 192, 224, 256]
+ def supports_head_size(cls, head_size: int) -> bool:
+ return head_size % 8 == 0 and head_size <= 256
@classmethod
def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
@@ -265,8 +288,8 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
self.dcp_world_size = 1
self.dcp_rank = 0
- self.dcp_kv_cache_interleave_size = (
- self.parallel_config.dcp_kv_cache_interleave_size
+ self.cp_kv_cache_interleave_size = (
+ self.parallel_config.cp_kv_cache_interleave_size
)
self.use_full_cuda_graph = (
@@ -388,7 +411,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
dcp_context_kv_lens_cpu,
self.dcp_world_size,
self.dcp_rank,
- self.dcp_kv_cache_interleave_size,
+ self.cp_kv_cache_interleave_size,
)
dcp_context_kv_lens = dcp_context_kv_lens_cpu.to(self.device)
max_dcp_context_kv_len = dcp_context_kv_lens.max().item()
diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py
index 4da1637d96eb6..8159f4096107f 100755
--- a/vllm/v1/attention/backends/flashinfer.py
+++ b/vllm/v1/attention/backends/flashinfer.py
@@ -16,7 +16,6 @@ from flashinfer import (
from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache
from flashinfer.prefill import trtllm_batch_context_with_kv_cache
from flashinfer.utils import FP4Tensor
-from typing_extensions import override
from vllm import envs
from vllm.attention.backends.abstract import (
@@ -275,10 +274,6 @@ class BatchDCPPrefillWrapper:
class FlashInferBackend(AttentionBackend):
accept_output_buffer: bool = True
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- # Note: Not sure for all platforms,
- # but on Blackwell, only support a page size of
- # 16, 32, 64
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [16, 32, 64]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"auto",
"fp8",
@@ -286,6 +281,12 @@ class FlashInferBackend(AttentionBackend):
"fp8_e5m2",
]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ # Note: Not sure for all platforms, but on Blackwell,
+ # only support a page size of 16, 32, 64.
+ return [16, 32, 64]
+
@staticmethod
def get_name() -> str:
return "FLASHINFER"
@@ -309,12 +310,20 @@ class FlashInferBackend(AttentionBackend):
return (num_blocks, 2, block_size, num_kv_heads, head_size)
@staticmethod
- def get_kv_cache_stride_order() -> tuple[int, ...]:
+ def get_kv_cache_stride_order(
+ include_num_layers_dimension: bool = False,
+ ) -> tuple[int, ...]:
# `stride_order` indicates the permutation that gets us from
# `get_kv_cache_shape` to the actual memory layout we want.
cache_layout = get_kv_cache_layout()
- if cache_layout == "NHD":
+ if cache_layout == "NHD" and include_num_layers_dimension:
+ # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
+ return (1, 0, 2, 3, 4, 5)
+ elif cache_layout == "NHD":
stride_order = (0, 1, 2, 3, 4)
+ elif cache_layout == "HND" and include_num_layers_dimension:
+ # (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
+ return (1, 2, 4, 0, 3, 5)
elif cache_layout == "HND":
stride_order = (0, 1, 3, 2, 4)
else:
@@ -558,7 +567,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
)
@classmethod
- @override
def get_cudagraph_support(
cls: type["FlashInferMetadataBuilder"],
vllm_config: VllmConfig,
@@ -585,6 +593,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
)
return self._workspace_buffer
+ def set_workspace_buffer(self, workspace_buffer: torch.Tensor):
+ self._workspace_buffer = workspace_buffer
+
def _get_prefill_wrapper(
self,
) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:
diff --git a/vllm/v1/attention/backends/mla/aiter_triton_mla.py b/vllm/v1/attention/backends/mla/aiter_triton_mla.py
new file mode 100644
index 0000000000000..8a92152a0ca53
--- /dev/null
+++ b/vllm/v1/attention/backends/mla/aiter_triton_mla.py
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.attention.backends.mla.common import MLACommonBackend
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
+ AiterMLAImpl,
+ AiterMLAMetadataBuilder,
+)
+
+
+class AiterTritonMLABackend(MLACommonBackend):
+ @staticmethod
+ def get_name() -> str:
+ return "AITER_TRITON_MLA"
+
+ @staticmethod
+ def get_impl_cls() -> type["AiterTritonMLAImpl"]:
+ return AiterTritonMLAImpl
+
+ @staticmethod
+ def get_builder_cls() -> type["AiterMLAMetadataBuilder"]:
+ return AiterMLAMetadataBuilder
+
+
+class AiterTritonMLAImpl(AiterMLAImpl):
+ def __init__(
+ self,
+ num_heads: int,
+ head_size: int,
+ scale: float,
+ num_kv_heads: int,
+ alibi_slopes: list[float] | None,
+ sliding_window: int | None,
+ kv_cache_dtype: str,
+ logits_soft_cap: float | None,
+ attn_type: str,
+ kv_sharing_target_layer_name: str | None,
+ # MLA Specific Arguments
+ **mla_args,
+ ) -> None:
+ super().__init__(
+ num_heads,
+ head_size,
+ scale,
+ num_kv_heads,
+ alibi_slopes,
+ sliding_window,
+ kv_cache_dtype,
+ logits_soft_cap,
+ attn_type,
+ kv_sharing_target_layer_name,
+ **mla_args,
+ )
+ from aiter.ops.triton.mha import flash_attn_varlen_func
+
+ self.flash_attn_varlen_func = flash_attn_varlen_func
+
+ def _flash_attn_varlen_diff_headdims(
+ self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
+ ):
+ result = self.flash_attn_varlen_func(
+ q,
+ k,
+ v,
+ softmax_scale=softmax_scale,
+ return_lse=return_softmax_lse,
+ **kwargs,
+ )
+ # Transpose the LSE if Triton MHA is used:
+ # (q.shape[0], num_q_heads) to (num_q_heads, q.shape[0])
+ if type(result) is tuple and return_softmax_lse:
+ output, lse = result
+ lse = lse.T.contiguous()
+ return (output, lse)
+ return result
diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py
index 2ccdd1f143ce8..43aef8a7cca91 100755
--- a/vllm/v1/attention/backends/mla/common.py
+++ b/vllm/v1/attention/backends/mla/common.py
@@ -308,6 +308,15 @@ class MLACommonBackend(AttentionBackend):
) -> tuple[int, ...]:
return (num_blocks, block_size, head_size)
+ @staticmethod
+ def get_kv_cache_stride_order(
+ include_num_layers_dimension: bool = False,
+ ) -> tuple[int, ...]:
+ # `stride_order` indicates the permutation that gets
+ # us from `get_kv_cache_shape` to the actual memory layout we want.
+ # (num_blocks, num_layers, block_size, head_size)
+ return (1, 0, 2, 3) if include_num_layers_dimension else (0, 1, 2)
+
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return [576]
@@ -536,7 +545,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
# DCP might not be initialized in testing
self.dcp_world_size = 1
self.dcp_rank = 0
- self.dcp_local_block_size = parallel_config.dcp_kv_cache_interleave_size
+ self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size
self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
# Don't try to access the runner on AMD
@@ -755,6 +764,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
seq_lens = common_attn_metadata.seq_lens
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
+ dcp_local_seq_lens_cpu = common_attn_metadata.dcp_local_seq_lens_cpu
query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
@@ -944,18 +954,20 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
decode_metadata = None
if num_decodes > 0:
+ dcp_tot_seq_lens_device = None
+ if self.dcp_world_size > 1:
+ dcp_tot_seq_lens_device = seq_lens[:num_decodes]
+ seq_lens_cpu = dcp_local_seq_lens_cpu
+ seq_lens = dcp_local_seq_lens
+
decode_metadata = self._build_decode(
block_table_tensor=block_table_tensor[:num_decodes, ...],
seq_lens_cpu=seq_lens_cpu[:num_decodes],
- seq_lens_device=dcp_local_seq_lens[:num_decodes]
- if self.dcp_world_size > 1 and dcp_local_seq_lens is not None
- else seq_lens[:num_decodes],
+ seq_lens_device=seq_lens[:num_decodes],
query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],
query_start_loc_device=query_start_loc[: num_decodes + 1],
num_decode_tokens=num_decode_tokens,
- dcp_tot_seq_lens_device=seq_lens[:num_decodes]
- if self.dcp_world_size > 1
- else None,
+ dcp_tot_seq_lens_device=dcp_tot_seq_lens_device,
)
attn_metadata = self.metadata_cls(
@@ -1286,8 +1298,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
get_current_vllm_config()
)
)
- self.dcp_kv_cache_interleave_size: int = (
- get_current_vllm_config().parallel_config.dcp_kv_cache_interleave_size
+ self.cp_kv_cache_interleave_size: int = (
+ get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
)
def _flash_attn_varlen_diff_headdims(
diff --git a/vllm/v1/attention/backends/mla/cutlass_mla.py b/vllm/v1/attention/backends/mla/cutlass_mla.py
index 60cb5022a55eb..5e3fbc0abf083 100644
--- a/vllm/v1/attention/backends/mla/cutlass_mla.py
+++ b/vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -36,13 +36,16 @@ class CutlassMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
class CutlassMLABackend(MLACommonBackend):
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [128]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"auto",
"fp8",
"fp8_e4m3",
]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [128]
+
@staticmethod
def get_name() -> str:
return "CUTLASS_MLA"
diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py
index 7794e89cc0a94..d369814c10b6f 100644
--- a/vllm/v1/attention/backends/mla/flashattn_mla.py
+++ b/vllm/v1/attention/backends/mla/flashattn_mla.py
@@ -41,9 +41,12 @@ logger = init_logger(__name__)
class FlashAttnMLABackend(MLACommonBackend):
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [MultipleOf(16)]
+
@staticmethod
def get_name() -> str:
return "FLASH_ATTN_MLA"
@@ -173,7 +176,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
) -> FlashAttnMLADecodeMetadata:
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
max_query_len = query_lens_cpu.max().item()
- max_seq_len = seq_lens_device.max().item()
+ max_seq_len = seq_lens_cpu.max().item()
# For Flash Attention MLA + full cudagraph
max_num_splits = 0
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py
index 52bb19e039e45..f02a4bb1ef35a 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py
@@ -35,13 +35,16 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
class FlashInferMLABackend(MLACommonBackend):
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [32, 64]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"auto",
"fp8",
"fp8_e4m3",
]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [32, 64]
+
@staticmethod
def get_name() -> str:
return "FLASHINFER_MLA"
diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py
index 3aab1f9bb7fb6..74a4cd8430250 100644
--- a/vllm/v1/attention/backends/mla/flashmla.py
+++ b/vllm/v1/attention/backends/mla/flashmla.py
@@ -39,13 +39,16 @@ logger = init_logger(__name__)
class FlashMLABackend(MLACommonBackend):
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"auto",
"fp8",
"fp8_e4m3",
]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [64]
+
@staticmethod
def get_name() -> str:
return "FLASHMLA"
diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py
index bb8d914d15719..1eee1d225293b 100644
--- a/vllm/v1/attention/backends/mla/flashmla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py
@@ -55,9 +55,12 @@ structured as:
class FlashMLASparseBackend(AttentionBackend):
accept_output_buffer: bool = True
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "fp8_ds_mla"]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [64]
+
@staticmethod
def get_name() -> str:
return "FLASHMLA_SPARSE"
@@ -168,7 +171,7 @@ def _convert_req_index_to_global_index_kernel(
inblock_off = tok % BLOCK_SIZE
# Guard block_table access
- valid_block = block_id < max_num_blocks_per_req
+ valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
base = tl.load(bt_ptr, mask=valid_block, other=0)
diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index 37aa5dad89a0e..77f1ba00d5b04 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -11,7 +11,8 @@ from vllm.attention.backends.abstract import (
)
from vllm.config import VllmConfig
from vllm.logger import init_logger
-from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata, is_deep_gemm_supported
from vllm.v1.attention.backends.utils import (
AttentionCGSupport,
AttentionMetadataBuilder,
@@ -23,7 +24,9 @@ logger = init_logger(__name__)
class DeepseekV32IndexerBackend(AttentionBackend):
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [1 if current_platform.is_rocm() else 64]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
@@ -45,7 +48,11 @@ class DeepseekV32IndexerBackend(AttentionBackend):
return (num_blocks, block_size, head_size)
@staticmethod
- def get_kv_cache_stride_order() -> tuple[int, ...]:
+ def get_kv_cache_stride_order(
+ include_num_layers_dimension: bool = False,
+ ) -> tuple[int, ...]:
+ if include_num_layers_dimension:
+ return (0, 1, 2, 3)
return (0, 1, 2)
@@ -328,10 +335,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
requires_padding = (decode_lens_cpu.max() > decode_lens_cpu.min()).item()
seq_lens = common_attn_metadata.seq_lens[:num_decodes]
-
- self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
- seq_lens, self.kv_cache_spec.block_size, self.num_sms
- )
+ if is_deep_gemm_supported():
+ self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
+ seq_lens, self.kv_cache_spec.block_size, self.num_sms
+ )
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
block_table=common_attn_metadata.block_table_tensor[:num_decodes, ...],
seq_lens=common_attn_metadata.seq_lens[:num_decodes],
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
index e1864526f02cc..56f9c7a281e7f 100644
--- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
@@ -7,9 +7,8 @@ from typing import ClassVar
import torch
from vllm._aiter_ops import rocm_aiter_ops
-from vllm.attention.backends.abstract import AttentionLayer
+from vllm.attention.backends.abstract import AttentionLayer, MultipleOf
from vllm.config import VllmConfig
-from vllm.utils.math_utils import cdiv
from vllm.v1.attention.backends.mla.common import (
MLACommonBackend,
MLACommonDecodeMetadata,
@@ -22,6 +21,10 @@ from vllm.v1.kv_cache_interface import AttentionSpec
class AiterMLABackend(MLACommonBackend):
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [1]
+
@staticmethod
def get_name() -> str:
return "ROCM_AITER_MLA"
@@ -71,9 +74,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
)
self.compilation_config = vllm_config.compilation_config
- max_num_pages_per_req = cdiv(
- vllm_config.model_config.max_model_len, self.kv_cache_spec.block_size
- )
+ # kernel block size is always 1.
+ max_num_pages_per_req = vllm_config.model_config.max_model_len
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
max_num_pages = max_num_reqs * max_num_pages_per_req
@@ -82,11 +84,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
# so we can only use the persistent buffer if a cudagraph is actually
# being used.
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
- self.block_table_remapping = torch.zeros(
- [max_num_reqs, max_num_pages_per_req * self.kv_cache_spec.block_size],
- dtype=torch.int32,
- device=device,
- )
self.paged_kv_indptr = torch.zeros(
max_num_reqs + 1, dtype=torch.int32, device=device
)
@@ -111,36 +108,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
num_decode_tokens: int,
dcp_tot_seq_lens_device: torch.Tensor | None,
) -> AiterMLADecodeMetadata:
- page_size = self.kv_cache_spec.block_size
+ # kernel block size is always 1, although the kv block size is not 1.
device = self.device
num_reqs = seq_lens_device.size(0)
- bs, _ = block_table_tensor.shape
- block_table_tensor = (
- block_table_tensor.unsqueeze(-1).expand(-1, -1, page_size) * page_size
- )
- block_table_tensor = (
- block_table_tensor
- + torch.arange(
- 0,
- page_size,
- device=block_table_tensor.device,
- dtype=block_table_tensor.dtype,
- )[None, None, :]
- )
- block_table_tensor = block_table_tensor.view(bs, -1)
- # after remapping, we assume the block size already equals to 1
-
- max_blk_size_per_req = block_table_tensor.shape[-1]
mask = torch.arange(
block_table_tensor.size(1), dtype=block_table_tensor.dtype, device=device
).unsqueeze(0) < seq_lens_device.unsqueeze(1)
paged_kv_indices = block_table_tensor[mask]
- paged_kv_last_page_len = seq_lens_device % page_size
- paged_kv_last_page_len = torch.where(
- paged_kv_last_page_len == 0, page_size, paged_kv_last_page_len
- )
+ paged_kv_last_page_len = torch.where(seq_lens_device == 0, 1, seq_lens_device)
paged_kv_indptr = torch.cat(
[
@@ -151,12 +128,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
num_actual_pages = paged_kv_indices.size(0)
- self.block_table_remapping[:num_reqs, :max_blk_size_per_req].copy_(
- block_table_tensor, non_blocking=True
- )
- block_table_tensor = self.block_table_remapping[
- :num_reqs, :max_blk_size_per_req
- ]
self.paged_kv_indices[:num_actual_pages].copy_(
paged_kv_indices, non_blocking=True
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
new file mode 100644
index 0000000000000..c0e7f0e380b98
--- /dev/null
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar, Optional
+
+import numpy as np
+import torch
+
+from vllm import _custom_ops as ops
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.attention.backends.abstract import (
+ AttentionBackend,
+ AttentionLayer,
+ AttentionMetadata,
+)
+from vllm.attention.backends.utils import get_mla_dims
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.v1.attention.backends.mla.common import (
+ MLACommonBaseImpl,
+)
+from vllm.v1.attention.backends.mla.flashmla_sparse import (
+ triton_convert_req_index_to_global_index,
+)
+from vllm.v1.attention.backends.utils import (
+ AttentionCGSupport,
+ AttentionMetadataBuilder,
+ CommonAttentionMetadata,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+if TYPE_CHECKING:
+ from vllm.model_executor.models.deepseek_v2 import Indexer
+logger = init_logger(__name__)
+
+
+class ROCMAiterMLASparseBackend(AttentionBackend):
+ accept_output_buffer: bool = True
+
+ @staticmethod
+ def get_name() -> str:
+ return "ROCM_AITER_MLA_SPARSE"
+
+ @staticmethod
+ def get_metadata_cls() -> type[AttentionMetadata]:
+ return ROCMAiterMLASparseMetadata
+
+ @staticmethod
+ def get_builder_cls() -> type["ROCMAiterMLASparseMetadataBuilder"]:
+ return ROCMAiterMLASparseMetadataBuilder
+
+ @staticmethod
+ def get_impl_cls() -> type["ROCMAiterMLASparseImpl"]:
+ return ROCMAiterMLASparseImpl
+
+ @staticmethod
+ def get_kv_cache_shape(
+ num_blocks: int,
+ block_size: int,
+ num_kv_heads: int, # assumed to be 1 for MLA
+ head_size: int,
+ cache_dtype_str: str = "auto",
+ ) -> tuple[int, ...]:
+ return (num_blocks, block_size, head_size)
+
+ @classmethod
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
+ return [torch.bfloat16]
+
+ @classmethod
+ def get_supported_head_sizes(cls) -> list[int]:
+ return [576]
+
+
+@dataclass
+class ROCMAiterMLASparseMetadata:
+ num_reqs: int
+ max_query_len: int
+ max_seq_len: int
+
+ num_actual_tokens: int # Number of tokens excluding padding.
+ query_start_loc: torch.Tensor
+ slot_mapping: torch.Tensor
+
+ block_table: torch.Tensor
+ req_id_per_token: torch.Tensor
+ block_size: int = 1
+ topk_tokens: int = 2048
+
+
+@dataclass
+class ROCMAiterMLASparseMetadataBuilder(
+ AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
+):
+ cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
+
+ def __init__(
+ self,
+ kv_cache_spec: AttentionSpec,
+ layer_names: list[str],
+ vllm_config: VllmConfig,
+ device: torch.device,
+ ):
+ self.kv_cache_spec = kv_cache_spec
+ self.model_config = vllm_config.model_config
+ parallel_config = vllm_config.parallel_config
+ self.device = device
+
+ self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
+ self.mla_dims = get_mla_dims(self.model_config)
+ self.topk_tokens = vllm_config.model_config.hf_config.index_topk
+ self.topk_tokens_tensor = torch.tensor(
+ [self.topk_tokens], device=device, dtype=torch.int32
+ )
+ self.max_model_len_tensor = torch.tensor(
+ [self.model_config.max_model_len], device=device, dtype=torch.int32
+ )
+ # this is ignored by `flash_mla_with_kvcache` if indices not None
+ self.dummy_block_table = torch.empty(
+ (1, 1), dtype=torch.int32, device=self.device
+ )
+
+ self.req_id_per_token_buffer = torch.empty(
+ (vllm_config.scheduler_config.max_num_batched_tokens,),
+ dtype=torch.int32,
+ device=device,
+ )
+
+ def build(
+ self,
+ common_prefix_len: int,
+ common_attn_metadata: CommonAttentionMetadata,
+ fast_build: bool = False,
+ ) -> ROCMAiterMLASparseMetadata:
+ num_tokens = common_attn_metadata.num_actual_tokens
+ starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
+ seg_lengths = np.diff(starts)
+ req_id_per_token = np.repeat(
+ np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
+ )
+ # Zero-fill for cudagraphs
+ self.req_id_per_token_buffer.fill_(0)
+ self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
+ torch.from_numpy(req_id_per_token), non_blocking=True
+ )
+ req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
+
+ metadata = ROCMAiterMLASparseMetadata(
+ num_reqs=common_attn_metadata.num_reqs,
+ max_query_len=common_attn_metadata.max_query_len,
+ max_seq_len=common_attn_metadata.max_seq_len,
+ num_actual_tokens=common_attn_metadata.num_actual_tokens,
+ query_start_loc=common_attn_metadata.query_start_loc,
+ slot_mapping=common_attn_metadata.slot_mapping,
+ block_table=common_attn_metadata.block_table_tensor,
+ req_id_per_token=req_id_per_token,
+ block_size=self.kv_cache_spec.block_size,
+ topk_tokens=self.topk_tokens,
+ )
+ return metadata
+
+
+# Take from
+# https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla_prefill.py#L72
+def reference_mla_sparse_prefill(
+ q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sm_scale: float, d_v: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+ import math
+
+ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor:
+ return torch.logsumexp(a * math.log(2), dim=dim) * math.log2(math.e)
+
+ skv = kv.shape[0]
+ sq = q.shape[0]
+ topk = indices.shape[-1]
+ dqk = q.shape[-1]
+ indices = indices[:, 0, :] # [s_q, topk]
+ invalid_indices_mask = (indices < 0) | (indices >= skv)
+ indices[invalid_indices_mask] = 0
+ qs = q # [s_q, h_q, d_qk]
+ kvs = kv[:, 0, :][indices].view(sq, topk, dqk) # [s_q, topk, d_qk]
+
+ attn_score = (qs @ kvs.transpose(1, 2)).float() # [s_q, h_q, topk]
+ attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf"))
+ attn_score *= sm_scale * math.log2(math.e)
+ lse = log2sumexp2(attn_score, dim=-1) # [s_q, h_q]
+ attn_score = torch.exp2(attn_score - lse.unsqueeze(-1)) # [s_q, h_q, topk]
+ result = attn_score.to(q.dtype) @ kvs[:, :, :d_v]
+ return (result, lse)
+
+
+class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
+ def __init__(
+ self,
+ num_heads: int,
+ head_size: int,
+ scale: float,
+ num_kv_heads: int,
+ alibi_slopes: list[float] | None,
+ sliding_window: int | None,
+ kv_cache_dtype: str,
+ logits_soft_cap: float | None,
+ attn_type: str,
+ kv_sharing_target_layer_name: str | None,
+ # MLA Specific Arguments
+ topk_indice_buffer: torch.Tensor | None = None,
+ indexer: Optional["Indexer"] = None,
+ **mla_args,
+ ) -> None:
+ super().__init__(
+ num_heads,
+ head_size,
+ scale,
+ num_kv_heads,
+ alibi_slopes,
+ sliding_window,
+ kv_cache_dtype,
+ logits_soft_cap,
+ attn_type,
+ kv_sharing_target_layer_name,
+ **mla_args,
+ )
+ self.softmax_scale = scale
+ assert indexer is not None
+ self.topk_indices_buffer = indexer.topk_indices_buffer
+ self.is_fp8bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
+
+ def _forward_bf16_kv(
+ self,
+ q: torch.Tensor,
+ kv_c_and_k_pe_cache: torch.Tensor,
+ topk_indices: torch.Tensor,
+ attn_metadata: ROCMAiterMLASparseMetadata,
+ ) -> torch.Tensor:
+ num_tokens = q.shape[0]
+ kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
+ -1, 1, kv_c_and_k_pe_cache.shape[-1]
+ )
+
+ topk_indices = topk_indices.view(num_tokens, 1, -1)
+ output = reference_mla_sparse_prefill(
+ q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale, 512
+ )[0]
+ return output[:, : self.num_heads, :]
+
+ def forward(
+ self,
+ layer: AttentionLayer,
+ q: torch.Tensor,
+ k_c_normed: torch.Tensor, # key in unified attn
+ k_pe: torch.Tensor, # value in unified attn
+ kv_cache: torch.Tensor,
+ attn_metadata: ROCMAiterMLASparseMetadata,
+ output: torch.Tensor | None = None,
+ output_scale: torch.Tensor | None = None,
+ output_block_scale: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
+ # MQA 576/512 approach for both prefill and decode
+
+ assert output is not None, "Output tensor must be provided."
+
+ if output_scale is not None or output_block_scale is not None:
+ raise NotImplementedError(
+ "fused output quantization is not yet supported for ROCMAiterMLASparse"
+ )
+
+ if attn_metadata is None:
+ # The zero fill is required when used with DP + EP
+ # to ensure all ranks within a DP group compute the
+ # same expert outputs.
+ return output.fill_(0)
+
+ num_actual_toks = attn_metadata.num_actual_tokens
+
+ # Inputs and outputs may be padded for CUDA graphs
+
+ q = q[:num_actual_toks, ...]
+ k_c_normed = k_c_normed[:num_actual_toks, ...]
+ k_pe = k_pe[:num_actual_toks, ...]
+
+ q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+ # Convert from (B, N, P) to (N, B, P)
+ q_nope = q_nope.transpose(0, 1)
+ if self.is_fp8bmm_enabled:
+ # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
+ ql_nope = rocm_aiter_ops.triton_fp8_bmm(
+ q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
+ )
+ else:
+ # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
+ ql_nope = torch.bmm(q_nope, self.W_UK_T)
+ # Convert from (N, B, L) to (B, N, L)
+ ql_nope = ql_nope.transpose(0, 1)
+
+ topk_indices = self.topk_indices_buffer[:num_actual_toks]
+
+ topk_indices_global = triton_convert_req_index_to_global_index(
+ attn_metadata.req_id_per_token,
+ attn_metadata.block_table,
+ topk_indices,
+ BLOCK_SIZE=attn_metadata.block_size,
+ NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
+ )
+
+ q = torch.cat([ql_nope, q_pe], dim=-1)
+
+ # write the latent and rope to kv cache
+ if kv_cache.numel() > 0:
+ ops.concat_and_cache_mla(
+ k_c_normed,
+ k_pe.squeeze(1),
+ kv_cache,
+ attn_metadata.slot_mapping.flatten(),
+ kv_cache_dtype=self.kv_cache_dtype,
+ scale=layer._k_scale,
+ )
+
+ attn_out = self._forward_bf16_kv(
+ q, kv_cache, topk_indices_global, attn_metadata
+ )
+
+ self._v_up_proj(attn_out, out=output[:num_actual_toks])
+ return output
diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py
index ea611848b0e81..ea911af3d19ce 100644
--- a/vllm/v1/attention/backends/rocm_aiter_fa.py
+++ b/vllm/v1/attention/backends/rocm_aiter_fa.py
@@ -447,7 +447,10 @@ class AiterFlashAttentionMetadataBuilder(
class AiterFlashAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
+
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [MultipleOf(16)]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
@@ -514,12 +517,9 @@ class AiterFlashAttentionImpl(AttentionImpl):
assert self.num_heads % self.num_kv_heads == 0
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
- if attn_type != AttentionType.DECODER:
+ if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
raise NotImplementedError(
- "Encoder self-attention and "
- "encoder/decoder cross-attention "
- "are not implemented for "
- "FlashAttentionImpl"
+ "Encoder self-attention is not implemented for FlashAttentionImpl"
)
def extend_forward(
@@ -675,7 +675,14 @@ class AiterFlashAttentionImpl(AttentionImpl):
# performance to make sure it does not introduce any overhead.
num_actual_tokens = attn_metadata.num_actual_tokens
key_cache, value_cache = kv_cache.unbind(0)
- if self.kv_sharing_target_layer_name is None:
+ # key and value may be None in the case of cross attention. They are
+ # calculated once based on the output from the encoder and then cached
+ # in KV cache.
+ if (
+ self.kv_sharing_target_layer_name is None
+ and key is not None
+ and value is not None
+ ):
# Reshape the input keys and values and store them in the cache.
# Skip this if sharing KV cache with an earlier attention layer.
# NOTE(woosuk): Here, key and value are padded while slot_mapping
@@ -701,8 +708,10 @@ class AiterFlashAttentionImpl(AttentionImpl):
# decode:extend:prefill
query = query[:num_actual_tokens]
- key = key[:num_actual_tokens]
- value = value[:num_actual_tokens]
+ if key is not None:
+ key = key[:num_actual_tokens]
+ if value is not None:
+ value = value[:num_actual_tokens]
output_actual_tokens = output[:num_actual_tokens]
diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
index b2639c0df0412..16fb52ab501c1 100644
--- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
@@ -142,7 +142,14 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
key_cache, value_cache = kv_cache.unbind(0)
- if self.kv_sharing_target_layer_name is None:
+ # key and value may be None in the case of cross attention. They are
+ # calculated once based on the output from the encoder and then cached
+ # in KV cache.
+ if (
+ self.kv_sharing_target_layer_name is None
+ and key is not None
+ and value is not None
+ ):
# Reshape the input keys and values and store them in the cache.
# Skip this if sharing KV cache with an earlier attention layer.
ops.reshape_and_cache_flash(
@@ -169,7 +176,10 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
max_seqlen_k = attn_metadata.max_seq_len
block_table = attn_metadata.block_table
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
+ descale_shape = (
+ cu_seqlens_q.shape[0] - 1,
+ key.shape[1] if key is not None else self.num_kv_heads,
+ )
self.unified_attention(
q=query[:num_actual_tokens],
diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py
index 6dfdfc19ccba1..868143cc192e7 100644
--- a/vllm/v1/attention/backends/rocm_attn.py
+++ b/vllm/v1/attention/backends/rocm_attn.py
@@ -238,12 +238,9 @@ class RocmAttentionImpl(AttentionImpl):
RocmAttentionBackend.validate_head_size(head_size)
- if attn_type != AttentionType.DECODER:
+ if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
raise NotImplementedError(
- "Encoder self-attention and "
- "encoder/decoder cross-attention "
- "are not implemented for "
- "RocmAttentionImpl"
+ "Encoder self-attention is not implemented for RocmAttentionImpl"
)
self.fp8_dtype = current_platform.fp8_dtype()
diff --git a/vllm/v1/attention/backends/tree_attn.py b/vllm/v1/attention/backends/tree_attn.py
index 1bf38ed225a4c..523f759e05a21 100644
--- a/vllm/v1/attention/backends/tree_attn.py
+++ b/vllm/v1/attention/backends/tree_attn.py
@@ -31,7 +31,10 @@ logger = init_logger(__name__)
class TreeAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
+
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [MultipleOf(16)]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py
index 889c79db18ef5..d051a89f03bb4 100644
--- a/vllm/v1/attention/backends/triton_attn.py
+++ b/vllm/v1/attention/backends/triton_attn.py
@@ -154,7 +154,6 @@ class TritonAttentionBackend(AttentionBackend):
torch.bfloat16,
torch.float32,
]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"auto",
"fp8",
@@ -162,6 +161,10 @@ class TritonAttentionBackend(AttentionBackend):
"fp8_e5m2",
]
+ @staticmethod
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+ return [MultipleOf(16)]
+
@staticmethod
def get_name() -> str:
return "TRITON_ATTN"
@@ -244,14 +247,11 @@ class TritonAttentionImpl(AttentionImpl):
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
- if attn_type != AttentionType.DECODER:
+ if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
raise NotImplementedError(
- "Encoder self-attention and "
- "encoder/decoder cross-attention "
- "are not implemented for "
- "TritonAttentionImpl"
+ "Encoder self-attention is not implemented for TritonAttentionImpl"
)
-
+ self.attn_type = attn_type
self.fp8_dtype = current_platform.fp8_dtype()
self.sinks = sinks
@@ -312,7 +312,11 @@ class TritonAttentionImpl(AttentionImpl):
num_actual_tokens = attn_metadata.num_actual_tokens
key_cache, value_cache = kv_cache.unbind(1)
- if self.kv_sharing_target_layer_name is None:
+ if (
+ self.kv_sharing_target_layer_name is None
+ and key is not None
+ and value is not None
+ ):
# Reshape the input keys and values and store them in the cache.
# Skip this if sharing KV cache with an earlier attention layer.
if self.kv_cache_dtype.startswith("fp8"):
@@ -346,7 +350,7 @@ class TritonAttentionImpl(AttentionImpl):
max_seqlen_k = attn_metadata.max_seq_len
block_table = attn_metadata.block_table
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])
unified_attention(
q=query[:num_actual_tokens],
diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py
index 578153cda7863..540a8e2b1d016 100644
--- a/vllm/v1/attention/backends/utils.py
+++ b/vllm/v1/attention/backends/utils.py
@@ -92,6 +92,7 @@ class CommonAttentionMetadata:
encoder_seq_lens: np.ndarray | None = None
dcp_local_seq_lens: torch.Tensor | None = None
+ dcp_local_seq_lens_cpu: torch.Tensor | None = None
"""Sequence lengths of the local rank in decode context parallelism world"""
@@ -1079,9 +1080,9 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
def get_dcp_local_seq_lens(
seq_lens: torch.Tensor,
- dcp_world_size: int = 1,
+ dcp_size: int = 1,
dcp_rank: int | None = None,
- dcp_kv_cache_interleave_size: int = 1,
+ cp_kv_cache_interleave_size: int = 1,
) -> torch.Tensor:
"""While using dcp, kv_cache size stored on each rank may be different,
use this function to calculate split decode seq_lens of each dcp rank.
@@ -1090,7 +1091,7 @@ def get_dcp_local_seq_lens(
num_requests = seq_lens.size(0)
if dcp_rank is None:
rank_offsets = (
- torch.arange(dcp_world_size, dtype=torch.int32)
+ torch.arange(dcp_size, dtype=torch.int32)
.unsqueeze(0)
.repeat(num_requests, 1)
)
@@ -1101,15 +1102,15 @@ def get_dcp_local_seq_lens(
)
base = (
seq_lens_tiled
- // dcp_kv_cache_interleave_size
- // dcp_world_size
- * dcp_kv_cache_interleave_size
+ // cp_kv_cache_interleave_size
+ // dcp_size
+ * cp_kv_cache_interleave_size
)
- remainder = seq_lens_tiled - base * dcp_world_size
+ remainder = seq_lens_tiled - base * dcp_size
remainder = torch.clip(
- remainder - rank_offsets * dcp_kv_cache_interleave_size,
+ remainder - rank_offsets * cp_kv_cache_interleave_size,
0,
- dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size,
)
dcp_local_seq_lens = base + remainder
return dcp_local_seq_lens.squeeze(1)
diff --git a/vllm/v1/attention/backends/xformers.py b/vllm/v1/attention/backends/xformers.py
deleted file mode 100644
index d15d79417cc61..0000000000000
--- a/vllm/v1/attention/backends/xformers.py
+++ /dev/null
@@ -1,417 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Attention layer with XFormersAttention."""
-
-from dataclasses import dataclass
-from typing import ClassVar, Optional
-
-import torch
-
-from vllm.attention.backends.abstract import (
- AttentionBackend,
- AttentionImpl,
- AttentionType,
- MultipleOf,
-)
-from vllm.attention.ops.triton_unified_attention import unified_attention
-from vllm.config import VllmConfig
-from vllm.logger import init_logger
-from vllm.v1.attention.backends.utils import (
- AttentionMetadataBuilder,
- CommonAttentionMetadata,
- split_decodes_and_prefills,
-)
-from vllm.v1.kv_cache_interface import AttentionSpec
-
-try:
- from xformers import ops as xops
- from xformers.ops.fmha.attn_bias import (
- AttentionBias,
- PagedBlockDiagonalCausalWithOffsetPaddedKeysMask,
- )
-
- XFORMERS_AVAILABLE = True
-except ImportError:
- XFORMERS_AVAILABLE = False
-
-from vllm import _custom_ops as ops
-
-logger = init_logger(__name__)
-
-
-class XFormersAttentionBackend(AttentionBackend):
- accept_output_buffer: bool = True
- supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
-
- @classmethod
- def get_supported_head_sizes(cls) -> list[int]:
- return [
- 32,
- 40,
- 48,
- 56,
- 64,
- 72,
- 80,
- 88,
- 96,
- 104,
- 112,
- 120,
- 128,
- 136,
- 144,
- 152,
- 160,
- 168,
- 176,
- 184,
- 192,
- 200,
- 208,
- 216,
- 224,
- 232,
- 240,
- 248,
- 256,
- ]
-
- @staticmethod
- def get_name() -> str:
- return "XFORMERS"
-
- @staticmethod
- def get_impl_cls() -> type["XFormersAttentionImpl"]:
- return XFormersAttentionImpl
-
- @staticmethod
- def get_kv_cache_shape(
- num_blocks: int,
- block_size: int,
- num_kv_heads: int,
- head_size: int,
- cache_dtype_str: str = "auto",
- ) -> tuple[int, ...]:
- if block_size % 16 != 0:
- raise ValueError("Block size must be a multiple of 16.")
- return (2, num_blocks, block_size, num_kv_heads, head_size)
-
- @staticmethod
- def get_builder_cls() -> type["XFormersAttentionMetadataBuilder"]:
- return XFormersAttentionMetadataBuilder
-
- @staticmethod
- def use_cascade_attention(*args, **kwargs) -> bool:
- return False
-
-
-@dataclass
-class XFormersAttentionMetadata:
- num_actual_tokens: int # Number of tokens excluding padding.
- max_query_len: int
- query_start_loc: torch.Tensor
- max_seq_len: int
- seq_lens: torch.Tensor
- block_table: torch.Tensor
- slot_mapping: torch.Tensor
-
- num_prefill_tokens: int = 0
- num_decode_tokens: int = 0
- num_prefills: int = 0
- num_decodes: int = 0
-
- # Biases for different attention types.
- attn_bias: Optional["AttentionBias"] = None
-
- # Self-attention prefill/decode metadata cache
- _cached_prefill_metadata: Optional["XFormersAttentionMetadata"] = None
- _cached_decode_metadata: Optional["XFormersAttentionMetadata"] = None
-
- @property
- def prefill_metadata(self) -> Optional["XFormersAttentionMetadata"]:
- if self.num_prefills == 0:
- return None
-
- if self._cached_prefill_metadata is not None:
- # Recover cached prefill-phase attention
- # metadata structure
- return self._cached_prefill_metadata
-
- q_start_loc = self.query_start_loc[self.num_decodes :]
- q_seqlens = torch.diff(q_start_loc)
- kv_seqlens = self.seq_lens[self.num_decodes :]
- # Construct & cache prefill-phase attention metadata structure
- self._cached_prefill_metadata = XFormersAttentionMetadata(
- num_actual_tokens=self.num_prefill_tokens,
- max_query_len=int(q_seqlens.max().item()),
- query_start_loc=q_start_loc - q_start_loc[0],
- max_seq_len=int(kv_seqlens.max().item()),
- seq_lens=kv_seqlens,
- block_table=self.block_table[self.num_decodes :],
- slot_mapping=self.slot_mapping[self.num_decode_tokens :],
- )
- return self._cached_prefill_metadata
-
- @property
- def decode_metadata(self) -> Optional["XFormersAttentionMetadata"]:
- if self.num_decode_tokens == 0:
- return None
-
- if self._cached_decode_metadata is not None:
- # Recover cached decode-phase attention
- # metadata structure
- return self._cached_decode_metadata
-
- q_start_loc = self.query_start_loc
- q_seqlens = torch.diff(q_start_loc)
- decode_kv_seqlens = self.seq_lens[: self.num_decodes]
- # Construct & cache decode-phase attention metadata structure
- self._cached_decode_metadata = XFormersAttentionMetadata(
- num_actual_tokens=self.num_decode_tokens,
- max_query_len=int(q_seqlens[: self.num_decodes].max().item()),
- query_start_loc=q_start_loc[: self.num_decodes + 1],
- max_seq_len=int(decode_kv_seqlens.max().item()),
- seq_lens=decode_kv_seqlens,
- block_table=self.block_table[: self.num_decodes],
- slot_mapping=self.slot_mapping[: self.num_decode_tokens],
- attn_bias=self.attn_bias,
- )
- return self._cached_decode_metadata
-
-
-class XFormersAttentionMetadataBuilder(
- AttentionMetadataBuilder[XFormersAttentionMetadata]
-):
- reorder_batch_threshold: int = 1
-
- def __init__(
- self,
- kv_cache_spec: AttentionSpec,
- layer_names: list[str],
- vllm_config: VllmConfig,
- device: torch.device,
- ):
- super().__init__(kv_cache_spec, layer_names, vllm_config, device)
-
- assert XFORMERS_AVAILABLE
- self.block_size = kv_cache_spec.block_size
- self._num_decodes = 0
- self._num_decode_tokens = 0
-
- def build(
- self,
- common_prefix_len: int,
- common_attn_metadata: CommonAttentionMetadata,
- fast_build: bool = False,
- ) -> XFormersAttentionMetadata:
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
- split_decodes_and_prefills(
- common_attn_metadata, decode_threshold=self.reorder_batch_threshold
- )
- )
-
- num_actual_tokens = common_attn_metadata.num_actual_tokens
- q_start_loc = common_attn_metadata.query_start_loc
- q_seqlens = torch.diff(q_start_loc)
- max_query_len = common_attn_metadata.max_query_len
- kv_seqlens = common_attn_metadata.seq_lens
- max_seq_len = common_attn_metadata.max_seq_len
- block_table = common_attn_metadata.block_table_tensor
- slot_mapping = common_attn_metadata.slot_mapping
-
- bias = None
- if num_decodes > 0:
- # Construct the decoder bias.
- decode_q_seqlens = q_seqlens[:num_decodes]
- decode_kv_seqlens = kv_seqlens[:num_decodes]
- bias = PagedBlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
- q_seqlen=decode_q_seqlens.tolist(),
- kv_seqlen=decode_kv_seqlens.tolist(),
- page_size=self.block_size,
- block_tables=block_table[:num_decodes],
- device=block_table.device,
- )
-
- return XFormersAttentionMetadata(
- num_actual_tokens=num_actual_tokens,
- num_prefill_tokens=num_prefill_tokens,
- num_decode_tokens=num_decode_tokens,
- num_prefills=num_prefills,
- num_decodes=num_decodes,
- max_query_len=max_query_len,
- query_start_loc=q_start_loc,
- max_seq_len=max_seq_len,
- seq_lens=kv_seqlens,
- block_table=block_table,
- slot_mapping=slot_mapping,
- attn_bias=bias,
- )
-
-
-class XFormersAttentionImpl(AttentionImpl):
- def __init__(
- self,
- num_heads: int,
- head_size: int,
- scale: float,
- num_kv_heads: int,
- alibi_slopes: list[float] | None,
- sliding_window: int | None,
- kv_cache_dtype: str,
- logits_soft_cap: float | None = None,
- attn_type: AttentionType = AttentionType.DECODER,
- kv_sharing_target_layer_name: str | None = None,
- ) -> None:
- if kv_sharing_target_layer_name is not None:
- raise NotImplementedError("KV sharing is not supported in V0.")
- if alibi_slopes is not None:
- raise NotImplementedError("XFormers does not support alibi slopes yet.")
- self.num_heads = num_heads
- self.head_size = head_size
- self.scale = float(scale)
- self.num_kv_heads = num_kv_heads
- self.num_queries_per_kv = self.num_heads // self.num_kv_heads
- self.kv_cache_dtype = kv_cache_dtype
- self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
- if alibi_slopes is not None:
- alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
- self.alibi_slopes = alibi_slopes
- if sliding_window is None:
- self.sliding_window = (-1, -1)
- else:
- self.sliding_window = (sliding_window - 1, 0)
- if logits_soft_cap is None:
- # Setting logits_soft_cap to 0 means no soft cap.
- logits_soft_cap = 0
- self.logits_soft_cap = logits_soft_cap
-
- if attn_type != AttentionType.DECODER:
- raise NotImplementedError(
- "Encoder self-attention and "
- "encoder/decoder cross-attention "
- "are not implemented for "
- "XFormersAttentionImpl."
- )
-
- def forward(
- self,
- layer: torch.nn.Module,
- query: torch.Tensor,
- key: torch.Tensor,
- value: torch.Tensor,
- kv_cache: torch.Tensor,
- attn_metadata: XFormersAttentionMetadata,
- output: torch.Tensor | None = None,
- output_scale: torch.Tensor | None = None,
- output_block_scale: torch.Tensor | None = None,
- ) -> torch.Tensor:
- """Forward pass with XFormers.
-
- Args:
- query: shape = [num_tokens, num_heads, head_size]
- key: shape = [num_tokens, num_kv_heads, head_size]
- value: shape = [num_tokens, num_kv_heads, head_size]
- kv_cache: shape =
- [2, num_blocks, block_size, num_kv_heads, head_size]
- attn_metadata: Metadata for attention.
- Returns:
- shape = [num_tokens, num_heads * head_size]
- """
- assert output is not None, "Output tensor must be provided."
-
- if output_scale is not None or output_block_scale is not None:
- raise NotImplementedError(
- "fused output quantization is not yet supported"
- " for XFormersAttentionImpl"
- )
-
- if attn_metadata is None:
- # Profiling run.
- return output.fill_(0)
-
- # Cache the input KVs.
- key_cache, value_cache = kv_cache.unbind(0)
- if self.kv_sharing_target_layer_name is None:
- # Reshape the input keys and values and store them in the cache.
- # Skip this if sharing KV cache with an earlier attention layer.
- # NOTE(woosuk): Here, key and value are padded while slot_mapping is
- # not padded. However, we don't need to do key[:num_actual_tokens]
- # and value[:num_actual_tokens] because the reshape_and_cache_flash
- # op uses the slot_mapping's shape to determine the number of
- # actual tokens.
- ops.reshape_and_cache_flash(
- key,
- value,
- key_cache,
- value_cache,
- attn_metadata.slot_mapping,
- self.kv_cache_dtype,
- layer._k_scale,
- layer._v_scale,
- )
-
- num_actual_tokens = attn_metadata.num_actual_tokens
- num_decode_tokens = attn_metadata.num_decode_tokens
- if prefill_meta := attn_metadata.prefill_metadata:
- descale_shape = (prefill_meta.query_start_loc.shape[0] - 1, key.shape[1])
- unified_attention(
- q=query[num_decode_tokens:num_actual_tokens],
- k=key_cache,
- v=value_cache,
- out=output[num_decode_tokens:num_actual_tokens],
- cu_seqlens_q=prefill_meta.query_start_loc,
- max_seqlen_q=prefill_meta.max_query_len,
- seqused_k=prefill_meta.seq_lens,
- max_seqlen_k=prefill_meta.max_seq_len,
- softmax_scale=self.scale,
- causal=True,
- alibi_slopes=self.alibi_slopes,
- window_size=self.sliding_window,
- block_table=prefill_meta.block_table,
- softcap=self.logits_soft_cap,
- q_descale=None, # Not supported
- k_descale=layer._k_scale.expand(descale_shape),
- v_descale=layer._v_scale.expand(descale_shape),
- )
-
- if decode_meta := attn_metadata.decode_metadata:
- # Query for decode. KV is not needed because it is already cached.
- decode_query = query[:num_decode_tokens]
- # Reshape query to [1, B_T, G, H, D].
- q = decode_query.view(
- 1, -1, self.num_kv_heads, self.num_queries_per_kv, self.head_size
- )
- # Reshape the k and v caches to [1, Bkv_T, G, H, D]
- cache_k = key_cache.view(
- 1, -1, self.num_kv_heads, 1, self.head_size
- ).expand(
- 1,
- -1,
- self.num_kv_heads,
- self.num_queries_per_kv,
- self.head_size,
- )
- cache_v = value_cache.view(
- 1, -1, self.num_kv_heads, 1, self.head_size
- ).expand(
- 1,
- -1,
- self.num_kv_heads,
- self.num_queries_per_kv,
- self.head_size,
- )
-
- attn_bias = decode_meta.attn_bias
- output[:num_decode_tokens] = xops.memory_efficient_attention_forward(
- q,
- cache_k,
- cache_v,
- attn_bias=attn_bias,
- p=0.0,
- scale=self.scale,
- ).view(decode_query.shape)
-
- # Reshape the output tensor.
- return output
diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py
index 137e5e0cdb6d2..1531b61f88fe2 100644
--- a/vllm/v1/core/kv_cache_coordinator.py
+++ b/vllm/v1/core/kv_cache_coordinator.py
@@ -27,6 +27,7 @@ class KVCacheCoordinator(ABC):
enable_caching: bool,
enable_kv_cache_events: bool,
dcp_world_size: int,
+ pcp_world_size: int,
):
self.kv_cache_config = kv_cache_config
self.max_model_len = max_model_len
@@ -44,6 +45,7 @@ class KVCacheCoordinator(ABC):
block_pool=self.block_pool,
kv_cache_group_id=i,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
)
@@ -210,6 +212,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
use_eagle: bool,
enable_kv_cache_events: bool,
dcp_world_size: int,
+ pcp_world_size: int,
):
super().__init__(
kv_cache_config,
@@ -218,6 +221,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
False,
enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
self.num_single_type_manager = len(self.single_type_managers)
@@ -250,6 +254,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
enable_caching: bool,
enable_kv_cache_events: bool,
dcp_world_size: int,
+ pcp_world_size: int,
):
super().__init__(
kv_cache_config,
@@ -258,12 +263,16 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
enable_caching,
enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
self.block_size = self.kv_cache_spec.block_size
self.dcp_world_size = dcp_world_size
+ self.pcp_world_size = pcp_world_size
if dcp_world_size > 1:
self.block_size *= dcp_world_size
+ if pcp_world_size > 1:
+ self.block_size *= pcp_world_size
assert len(self.kv_cache_config.kv_cache_groups) == 1, (
"UnitaryKVCacheCoordinator assumes only one kv cache group"
)
@@ -281,6 +290,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
kv_cache_spec=self.kv_cache_spec,
use_eagle=self.use_eagle,
dcp_world_size=self.dcp_world_size,
+ pcp_world_size=self.pcp_world_size,
)
return hit_blocks, len(hit_blocks[0]) * self.block_size
@@ -302,6 +312,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
enable_caching: bool,
enable_kv_cache_events: bool,
dcp_world_size: int,
+ pcp_world_size: int,
):
super().__init__(
kv_cache_config,
@@ -310,8 +321,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
enable_caching,
enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
assert dcp_world_size == 1, "DCP not support hybrid attn now."
+ assert pcp_world_size == 1, "PCP not support hybrid attn now."
self.verify_and_split_kv_cache_groups()
def verify_and_split_kv_cache_groups(self) -> None:
@@ -452,6 +465,7 @@ def get_kv_cache_coordinator(
enable_caching: bool,
enable_kv_cache_events: bool,
dcp_world_size: int,
+ pcp_world_size: int,
) -> KVCacheCoordinator:
if not enable_caching:
return KVCacheCoordinatorNoPrefixCache(
@@ -460,6 +474,7 @@ def get_kv_cache_coordinator(
use_eagle,
enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
if len(kv_cache_config.kv_cache_groups) == 1:
return UnitaryKVCacheCoordinator(
@@ -469,6 +484,7 @@ def get_kv_cache_coordinator(
enable_caching,
enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
return HybridKVCacheCoordinator(
kv_cache_config,
@@ -477,4 +493,5 @@ def get_kv_cache_coordinator(
enable_caching,
enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 7f405fc248ac2..2012c3fef88bc 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -100,6 +100,7 @@ class KVCacheManager:
log_stats: bool = False,
enable_kv_cache_events: bool = False,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> None:
self.max_model_len = max_model_len
@@ -124,12 +125,9 @@ class KVCacheManager:
0
].kv_cache_spec.block_size
- if dcp_world_size > 1:
+ if dcp_world_size * pcp_world_size > 1:
assert len(kv_cache_config.kv_cache_groups) == 1
- # Note(hc): need revisit. When both DCP and any future
- # PCP are enabled, the block_size may need to be scaled
- # by a factor of dcp_size × pcp_size?
- self.block_size *= dcp_world_size
+ self.block_size *= dcp_world_size * pcp_world_size
self.coordinator = get_kv_cache_coordinator(
kv_cache_config=kv_cache_config,
@@ -138,6 +136,7 @@ class KVCacheManager:
enable_caching=self.enable_caching,
enable_kv_cache_events=enable_kv_cache_events,
dcp_world_size=dcp_world_size,
+ pcp_world_size=pcp_world_size,
)
self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
self.block_pool = self.coordinator.block_pool
diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py
index 6e026215d4022..b18ba8e8b2c7b 100644
--- a/vllm/v1/core/kv_cache_utils.py
+++ b/vllm/v1/core/kv_cache_utils.py
@@ -1219,11 +1219,16 @@ def _report_kv_cache_config(
// len(kv_cache_config.kv_cache_groups)
* min_block_size
)
- if vllm_config.parallel_config.decode_context_parallel_size > 1:
- num_tokens *= vllm_config.parallel_config.decode_context_parallel_size
+ dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+ pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
+ if pcp_size * dcp_size > 1:
+ num_tokens *= pcp_size * dcp_size
logger.info(
- "Multiplying the GPU KV cache size by the dcp_world_size %d.",
- vllm_config.parallel_config.decode_context_parallel_size,
+ "Multiplying the GPU KV cache size by the cp_world_size %d "
+ "(pcp_world_size %d * dcp_world_size %d).",
+ pcp_size * dcp_size,
+ pcp_size,
+ dcp_size,
)
num_tokens_str = f"{num_tokens:,}"
logger.info_once("GPU KV cache size: %s tokens", num_tokens_str, scope="local")
@@ -1231,10 +1236,11 @@ def _report_kv_cache_config(
max_concurrency = get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config
)
- logger.info(
+ logger.info_once(
"Maximum concurrency for %s tokens per request: %.2fx",
max_model_len_str,
max_concurrency,
+ scope="local",
)
diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py
index 20fdb3446404b..7902513dce49a 100644
--- a/vllm/v1/core/sched/output.py
+++ b/vllm/v1/core/sched/output.py
@@ -44,11 +44,15 @@ class NewRequestData:
lora_request: LoRARequest | None
prompt_embeds: "torch.Tensor | None" = None
+ # Only used for v2 model runner.
+ prefill_token_ids: list[int] | None = None
+
@classmethod
def from_request(
cls,
request: Request,
block_ids: tuple[list[int], ...],
+ prefill_token_ids: list[int] | None = None,
) -> "NewRequestData":
return cls(
req_id=request.request_id,
@@ -60,6 +64,7 @@ class NewRequestData:
num_computed_tokens=request.num_computed_tokens,
lora_request=request.lora_request,
prompt_embeds=request.prompt_embeds,
+ prefill_token_ids=prefill_token_ids,
)
def __repr__(self) -> str:
@@ -68,6 +73,7 @@ class NewRequestData:
f"NewRequestData("
f"req_id={self.req_id},"
f"prompt_token_ids={self.prompt_token_ids},"
+ f"prefill_token_ids={self.prefill_token_ids},"
f"mm_features={self.mm_features},"
f"sampling_params={self.sampling_params},"
f"block_ids={self.block_ids},"
@@ -183,6 +189,10 @@ class SchedulerOutput:
# freed from the encoder cache.
free_encoder_mm_hashes: list[str]
+ # Request IDs that are preempted in this step.
+ # Only used for v2 model runner.
+ preempted_req_ids: set[str] | None = None
+
# Whether the scheduled requests have all the output tokens they
# need to perform grammar bitmask computation.
pending_structured_output_tokens: bool = False
@@ -193,6 +203,20 @@ class SchedulerOutput:
# EC Cache Connector metadata
ec_connector_metadata: ECConnectorMetadata | None = None
+ @classmethod
+ def make_empty(cls) -> "SchedulerOutput":
+ return cls(
+ scheduled_new_reqs=[],
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
+ num_scheduled_tokens={},
+ total_num_scheduled_tokens=0,
+ scheduled_spec_decode_tokens={},
+ scheduled_encoder_inputs={},
+ num_common_prefix_blocks=[],
+ finished_req_ids=set(),
+ free_encoder_mm_hashes=[],
+ )
+
@dataclass
class GrammarOutput:
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 4323141c435b7..a7ec0de372631 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -6,6 +6,7 @@ from collections import defaultdict
from collections.abc import Iterable
from typing import Any
+from vllm import envs
from vllm.config import VllmConfig
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorMetadata,
@@ -121,6 +122,7 @@ class Scheduler(SchedulerInterface):
self.block_size = block_size
self.dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
+ self.pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
# req_id -> Request
self.requests: dict[str, Request] = {}
@@ -178,13 +180,15 @@ class Scheduler(SchedulerInterface):
self.kv_cache_manager = KVCacheManager(
kv_cache_config=kv_cache_config,
max_model_len=self.max_model_len,
- enable_caching=bool(self.cache_config.enable_prefix_caching),
+ enable_caching=self.cache_config.enable_prefix_caching,
use_eagle=self.use_eagle,
log_stats=self.log_stats,
enable_kv_cache_events=self.enable_kv_cache_events,
dcp_world_size=self.dcp_world_size,
+ pcp_world_size=self.pcp_world_size,
)
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
+ self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
def schedule(self) -> SchedulerOutput:
# NOTE(woosuk) on the scheduling algorithm:
@@ -466,6 +470,7 @@ class Scheduler(SchedulerInterface):
skipped_waiting_requests.prepend_request(request)
continue
+ request.num_external_computed_tokens = ext_tokens
num_external_computed_tokens = ext_tokens
# Total computed tokens (local + external).
@@ -572,9 +577,6 @@ class Scheduler(SchedulerInterface):
new_computed_blocks + new_blocks,
num_external_computed_tokens,
)
- self._update_connector_prefix_cache_stats(
- request, num_external_computed_tokens
- )
# Request was already popped from self.waiting
# unless it was re-added above due to new_blocks being None.
@@ -586,6 +588,8 @@ class Scheduler(SchedulerInterface):
request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
continue
+ self._update_connector_prefix_cache_stats(request)
+
req_index += 1
self.running.append(request)
if self.log_stats:
@@ -656,12 +660,25 @@ class Scheduler(SchedulerInterface):
)
# Construct the scheduler output.
- new_reqs_data = [
- NewRequestData.from_request(
- req, req_to_new_blocks[req.request_id].get_block_ids()
- )
- for req in scheduled_new_reqs
- ]
+ if self.use_v2_model_runner:
+ scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs
+ scheduled_resumed_reqs = []
+ new_reqs_data = [
+ NewRequestData.from_request(
+ req,
+ req_to_new_blocks[req.request_id].get_block_ids(),
+ req._all_token_ids,
+ )
+ for req in scheduled_new_reqs
+ ]
+ else:
+ new_reqs_data = [
+ NewRequestData.from_request(
+ req, req_to_new_blocks[req.request_id].get_block_ids()
+ )
+ for req in scheduled_new_reqs
+ ]
+
with record_function_or_nullcontext("schedule: make_cached_request_data"):
cached_reqs_data = self._make_cached_request_data(
scheduled_running_reqs,
@@ -683,6 +700,7 @@ class Scheduler(SchedulerInterface):
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
scheduled_encoder_inputs=scheduled_encoder_inputs,
num_common_prefix_blocks=num_common_prefix_blocks,
+ preempted_req_ids={req.request_id for req in preempted_reqs},
# finished_req_ids is an existing state in the scheduler,
# instead of being newly scheduled in this step.
# It contains the request IDs that are finished in between
@@ -1011,8 +1029,8 @@ class Scheduler(SchedulerInterface):
continue
req_index = model_runner_output.req_id_to_index[req_id]
- generated_token_ids: list[int] = (
- sampled_token_ids[req_index].tolist() if sampled_token_ids else []
+ generated_token_ids = (
+ sampled_token_ids[req_index] if sampled_token_ids else []
)
scheduled_spec_token_ids = (
@@ -1362,15 +1380,13 @@ class Scheduler(SchedulerInterface):
# KV Connector Related Methods
########################################################################
- def _update_connector_prefix_cache_stats(
- self, request: Request, num_external_tokens: int
- ) -> None:
+ def _update_connector_prefix_cache_stats(self, request: Request) -> None:
if self.connector_prefix_cache_stats is None:
return
self.connector_prefix_cache_stats.record(
num_tokens=request.num_tokens,
- num_hits=num_external_tokens,
+ num_hits=request.num_external_computed_tokens,
preempted=request.num_preemptions > 0,
)
@@ -1553,9 +1569,11 @@ class Scheduler(SchedulerInterface):
marked_invalid_block = True
# Truncate the computed tokens at the first failed block
request.num_computed_tokens = idx * self.block_size
- total_affected_tokens += (
+ num_affected_tokens = (
req_num_computed_tokens - request.num_computed_tokens
)
+ total_affected_tokens += num_affected_tokens
+ request.num_external_computed_tokens -= num_affected_tokens
if is_affected:
if not marked_invalid_block:
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
index 14ac83028ee44..d90ec550f7666 100644
--- a/vllm/v1/core/single_type_kv_cache_manager.py
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
@@ -32,6 +32,7 @@ class SingleTypeKVCacheManager(ABC):
block_pool: BlockPool,
kv_cache_group_id: int,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> None:
"""
Initializes the SingleTypeKVCacheManager.
@@ -42,8 +43,9 @@ class SingleTypeKVCacheManager(ABC):
"""
self.block_size = kv_cache_spec.block_size
self.dcp_world_size = dcp_world_size
- if self.dcp_world_size > 1:
- self.block_size *= dcp_world_size
+ self.pcp_world_size = pcp_world_size
+ if dcp_world_size * pcp_world_size > 1:
+ self.block_size *= dcp_world_size * pcp_world_size
self.kv_cache_spec = kv_cache_spec
self.block_pool = block_pool
@@ -212,6 +214,7 @@ class SingleTypeKVCacheManager(ABC):
kv_cache_spec: KVCacheSpec,
use_eagle: bool,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> tuple[list[KVCacheBlock], ...]:
"""
Get the longest cache hit prefix of the blocks that is not longer than
@@ -303,6 +306,7 @@ class FullAttentionManager(SingleTypeKVCacheManager):
kv_cache_spec: KVCacheSpec,
use_eagle: bool,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> tuple[list[KVCacheBlock], ...]:
assert isinstance(
kv_cache_spec, (FullAttentionSpec, ChunkedLocalAttentionSpec)
@@ -314,8 +318,8 @@ class FullAttentionManager(SingleTypeKVCacheManager):
[] for _ in range(len(kv_cache_group_ids))
)
block_size = kv_cache_spec.block_size
- if dcp_world_size > 1:
- block_size *= dcp_world_size
+ if dcp_world_size * pcp_world_size > 1:
+ block_size *= dcp_world_size * pcp_world_size
max_num_blocks = max_length // block_size
for block_hash in itertools.islice(block_hashes, max_num_blocks):
# block_hashes is a chain of block hashes. If a block hash is not
@@ -362,11 +366,13 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
kv_cache_spec: KVCacheSpec,
use_eagle: bool,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> tuple[list[KVCacheBlock], ...]:
assert isinstance(kv_cache_spec, SlidingWindowSpec), (
"SlidingWindowManager can only be used for sliding window groups"
)
assert dcp_world_size == 1, "DCP not support sliding window attn now."
+ assert pcp_world_size == 1, "PCP not support sliding window attn now."
# The number of contiguous blocks needed for prefix cache hit.
# -1 since the input token itself is also included in the window
@@ -476,6 +482,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
kv_cache_spec: KVCacheSpec,
use_eagle: bool,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> tuple[list[KVCacheBlock], ...]:
"""
For chunked local attention, we need to find the longest cache hit
@@ -516,6 +523,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
"Hybrid KV cache is not supported for " + "eagle + chunked local attention."
)
assert dcp_world_size == 1, "DCP not support chunked local attn now."
+ assert pcp_world_size == 1, "PCP not support chunked local attn now."
max_num_blocks = max_length // kv_cache_spec.block_size
if max_length > 0:
local_attention_start_idx = (
@@ -611,11 +619,13 @@ class MambaManager(SingleTypeKVCacheManager):
kv_cache_spec: KVCacheSpec,
use_eagle: bool,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> tuple[list[KVCacheBlock], ...]:
assert isinstance(kv_cache_spec, MambaSpec), (
"MambaManager can only be used for mamba groups"
)
assert dcp_world_size == 1, "DCP not support mamba now."
+ assert pcp_world_size == 1, "PCP not support mamba now."
computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
[] for _ in range(len(kv_cache_group_ids))
)
@@ -705,6 +715,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
kv_cache_spec: KVCacheSpec,
use_eagle: bool,
dcp_world_size: int = 1,
+ pcp_world_size: int = 1,
) -> tuple[list[KVCacheBlock], ...]:
assert isinstance(kv_cache_spec, CrossAttentionSpec), (
"CrossAttentionManager can only be used for cross-attention groups"
diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py
index c160c7cbcab4a..c64b3cccfc652 100644
--- a/vllm/v1/engine/async_llm.py
+++ b/vllm/v1/engine/async_llm.py
@@ -152,6 +152,10 @@ class AsyncLLM(EngineClient):
)
self.logger_manager.log_engine_initialized()
+ # Pause / resume state for async RL workflows.
+ self._pause_cond = asyncio.Condition()
+ self._paused = False
+
self.output_handler: asyncio.Task | None = None
try:
# Start output handler eagerly if we are in the asyncio eventloop.
@@ -160,11 +164,23 @@ class AsyncLLM(EngineClient):
except RuntimeError:
pass
- if envs.VLLM_TORCH_PROFILER_DIR:
+ if (
+ envs.VLLM_TORCH_PROFILER_DIR
+ and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM
+ ):
logger.info(
"Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501
envs.VLLM_TORCH_PROFILER_DIR,
)
+ if envs.VLLM_PROFILER_MAX_ITERS > 0 or envs.VLLM_PROFILER_DELAY_ITERS > 0:
+ logger.warning_once(
+ "Torch profiler received max_iters or delay_iters setting. These "
+ "are not compatible with the AsyncLLM profiler and will be ignored "
+ "for the AsyncLLM process. Engine process profiling will still "
+ "respect these settings. Consider setting "
+ "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM=1 to disable "
+ "AsyncLLM profiling."
+ )
worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
self.profiler = torch.profiler.profile(
activities=[
@@ -392,6 +408,10 @@ class AsyncLLM(EngineClient):
# to handle startup failure gracefully in the OpenAI server.
self._run_output_handler()
+ # Wait until generation is resumed if the engine is paused.
+ async with self._pause_cond:
+ await self._pause_cond.wait_for(lambda: not self._paused)
+
if tokenization_kwargs is None:
tokenization_kwargs = {}
truncate_prompt_tokens = sampling_params.truncate_prompt_tokens
@@ -539,6 +559,58 @@ class AsyncLLM(EngineClient):
if self.log_requests:
logger.info("Aborted request(s) %s.", ",".join(request_ids))
+ async def pause_generation(
+ self,
+ *,
+ wait_for_inflight_requests: bool = False,
+ clear_cache: bool = True,
+ ) -> None:
+ """
+ Pause generation to allow model weight updates.
+
+ New generation/encoding requests are blocked until resume.
+
+ Args:
+ wait_for_inflight_requests: When ``True`` waits for in-flight
+ requests to finish before pausing. When ``False`` (default),
+ immediately aborts any in-flight requests.
+ clear_cache: Whether to clear KV cache and prefix cache after
+ draining. Set to ``False`` to preserve cache for faster resume.
+ Default is ``True`` (clear caches).
+ """
+
+ async with self._pause_cond:
+ if self._paused:
+ return
+ self._paused = True
+
+ if not wait_for_inflight_requests:
+ request_ids = list(self.output_processor.request_states.keys())
+ if request_ids:
+ await self.abort(request_ids)
+
+ # Wait for running requests to drain before clearing cache.
+ if self.output_processor.has_unfinished_requests():
+ await self.output_processor.wait_for_requests_to_drain()
+
+ # Clear cache
+ if clear_cache:
+ await self.reset_prefix_cache()
+ await self.reset_mm_cache()
+
+ async def resume_generation(self) -> None:
+ """Resume generation after :meth:`pause_generation`."""
+
+ async with self._pause_cond:
+ self._paused = False
+ self._pause_cond.notify_all() # Wake up all waiting requests
+
+ async def is_paused(self) -> bool:
+ """Return whether the engine is currently paused."""
+
+ async with self._pause_cond:
+ return self._paused
+
async def encode(
self,
prompt: PromptType,
@@ -570,6 +642,10 @@ class AsyncLLM(EngineClient):
# to handle startup failure gracefully in the OpenAI server.
self._run_output_handler()
+ # Respect pause state before accepting new requests.
+ async with self._pause_cond:
+ await self._pause_cond.wait_for(lambda: not self._paused)
+
if tokenization_kwargs is None:
tokenization_kwargs = {}
_validate_truncation_size(
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
index 97286c6e2e5e4..8657a95b5e6e7 100644
--- a/vllm/v1/engine/core.py
+++ b/vllm/v1/engine/core.py
@@ -128,6 +128,7 @@ class EngineCore:
scheduler_block_size = (
vllm_config.cache_config.block_size
* vllm_config.parallel_config.decode_context_parallel_size
+ * vllm_config.parallel_config.prefill_context_parallel_size
)
self.scheduler: SchedulerInterface = Scheduler(
@@ -180,10 +181,11 @@ class EngineCore:
logger.info("Batch queue is enabled with size %d", self.batch_queue_size)
self.batch_queue = deque(maxlen=self.batch_queue_size)
- self.ec_producer = (
+ self.is_ec_producer = (
vllm_config.ec_transfer_config is not None
and vllm_config.ec_transfer_config.is_ec_producer
)
+ self.is_pooling_model = vllm_config.model_config.runner_type == "pooling"
self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None
if vllm_config.cache_config.enable_prefix_caching or kv_connector is not None:
@@ -204,6 +206,8 @@ class EngineCore:
# Mark the startup heap as static so that it's ignored by GC.
# Reduces pause times of oldest generation collections.
freeze_gc_heap()
+ # If enable, attach GC debugger after static variable freeze.
+ maybe_attach_gc_debug_callback()
def _initialize_kv_caches(
self, vllm_config: VllmConfig
@@ -389,10 +393,10 @@ class EngineCore:
exec_future = self.model_executor.execute_model(
scheduler_output, non_block=True
)
- if not self.ec_producer:
+ if not self.is_ec_producer:
model_executed = scheduler_output.total_num_scheduled_tokens > 0
- if not model_executed:
+ if self.is_pooling_model or not model_executed:
# No sampling required (no requests scheduled).
future = cast(Future[ModelRunnerOutput], exec_future)
else:
@@ -643,9 +647,6 @@ class EngineCoreProc(EngineCore):
assert addresses.coordinator_input is not None
logger.info("Waiting for READY message from DP Coordinator...")
- # If enable, attach GC debugger after static variable freeze.
- maybe_attach_gc_debug_callback()
-
# Enable environment variable cache (e.g. assume no more
# environment variable overrides after this point)
enable_envs_cache()
diff --git a/vllm/v1/engine/logprobs.py b/vllm/v1/engine/logprobs.py
index b618d23472651..63064a2c65d67 100644
--- a/vllm/v1/engine/logprobs.py
+++ b/vllm/v1/engine/logprobs.py
@@ -43,15 +43,22 @@ class LogprobsProcessor:
tokenizer: AnyTokenizer | None,
request: EngineCoreRequest,
) -> "LogprobsProcessor":
- assert request.sampling_params is not None
- num_logprobs = request.sampling_params.logprobs
- num_prompt_logprobs = request.sampling_params.prompt_logprobs
+ sampling_params = request.sampling_params
+ assert sampling_params is not None
+ num_logprobs = sampling_params.logprobs
+ num_prompt_logprobs = sampling_params.prompt_logprobs
return cls(
tokenizer=tokenizer,
cumulative_logprob=(None if num_logprobs is None else 0.0),
- logprobs=(None if num_logprobs is None else create_sample_logprobs()),
+ logprobs=(
+ None
+ if num_logprobs is None
+ else create_sample_logprobs(sampling_params.flat_logprobs)
+ ),
prompt_logprobs=(
- None if num_prompt_logprobs is None else create_prompt_logprobs()
+ None
+ if num_prompt_logprobs is None
+ else create_prompt_logprobs(sampling_params.flat_logprobs)
),
num_prompt_logprobs=num_prompt_logprobs,
num_logprobs=num_logprobs,
diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py
index bdbbfe2595f81..0453c4a77f0cd 100644
--- a/vllm/v1/engine/output_processor.py
+++ b/vllm/v1/engine/output_processor.py
@@ -350,6 +350,8 @@ class OutputProcessor:
self.parent_requests: dict[str, ParentRequest] = {}
self.lora_states = LoRARequestStates(log_stats)
self.tracer: Tracer | None = None
+ self._requests_drained = asyncio.Event()
+ self._requests_drained.set()
def get_num_unfinished_requests(self):
return len(self.request_states)
@@ -357,6 +359,11 @@ class OutputProcessor:
def has_unfinished_requests(self) -> bool:
return len(self.request_states) > 0
+ async def wait_for_requests_to_drain(self) -> None:
+ if not self.request_states:
+ return
+ await self._requests_drained.wait()
+
def propagate_error(self, e: Exception):
"""Propagate error to all generate() tasks."""
@@ -396,6 +403,8 @@ class OutputProcessor:
child_reqs = self.abort_requests(child_reqs)
request_ids_to_abort.extend(child_reqs)
self.parent_requests.pop(request_id, None)
+ if not self.request_states:
+ self._requests_drained.set()
return request_ids_to_abort
def add_request(
@@ -420,6 +429,8 @@ class OutputProcessor:
log_stats=self.log_stats,
stream_interval=self.stream_interval,
)
+ if self._requests_drained.is_set():
+ self._requests_drained.clear()
self.request_states[request_id] = req_state
if parent_req:
self.parent_requests[parent_req.request_id] = parent_req
@@ -511,6 +522,8 @@ class OutputProcessor:
parent_req = req_state.parent_req
if parent_req and not parent_req.child_requests:
self.parent_requests.pop(parent_req.request_id, None)
+ if not self.request_states:
+ self._requests_drained.set()
if not engine_core_output.finished:
# If req not finished in EngineCore, but Detokenizer
# detected stop string, abort needed in EngineCore.
diff --git a/vllm/v1/engine/processor.py b/vllm/v1/engine/processor.py
index 4cb911d8e22b7..af4f0e410e253 100644
--- a/vllm/v1/engine/processor.py
+++ b/vllm/v1/engine/processor.py
@@ -20,6 +20,7 @@ from vllm.multimodal.utils import argsort_mm_positions
from vllm.pooling_params import PoolingParams
from vllm.sampling_params import SamplingParams
from vllm.transformers_utils.tokenizer import AnyTokenizer
+from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
from vllm.utils import length_from_prompt_token_ids_or_embeds
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.metrics.stats import MultiModalCacheStats
@@ -142,9 +143,6 @@ class Processor:
self,
params: SamplingParams,
) -> None:
- # Best of not yet supported.
- if params.best_of is not None and params.best_of > 1:
- raise ValueError("vLLM V1 does not yet support best_of.")
# Logits processors not supported.
if params.logits_processors:
raise ValueError(
@@ -303,12 +301,24 @@ class Processor:
# allows <|special_token|> and similar, see
# https://github.com/guidance-ai/llguidance/blob/main/docs/syntax.md#special-tokens
# Without tokenizer these are disallowed in grammars.
+ if isinstance(self.tokenizer, MistralTokenizer):
+ raise ValueError(
+ "Mistral tokenizer is not supported for the 'guidance' "
+ "structured output backend. Please use ['xgrammar', 'outlines'] "
+ "backends or tokenizer_mode='hf' instead."
+ )
validate_guidance_grammar(params, tokenizer=None)
elif backend == "outlines":
# outlines backend
validate_structured_output_request_outlines(params)
elif backend == "lm-format-enforcer":
# lm format enforcer backend
+ if isinstance(self.tokenizer, MistralTokenizer):
+ raise ValueError(
+ "Mistral tokenizer is not supported for the 'lm-format-enforcer' "
+ "structured output backend. Please use ['xgrammar', 'outlines'] "
+ "backends or tokenizer_mode='hf' instead."
+ )
validate_structured_output_request_lm_format_enforcer(params)
else:
# NOTE: backend must be "auto" here, because we have
@@ -323,9 +333,15 @@ class Processor:
except ValueError:
# The request either failed validation
# or includes some jsonschema feature(s) that
- # are not supported in xgrammar. Fall back to guidance.
- validate_guidance_grammar(params, tokenizer=None)
- params.structured_outputs._backend = "guidance"
+ # are not supported in xgrammar.
+ if isinstance(self.tokenizer, MistralTokenizer):
+ # Fall back to outlines if the tokenizer is Mistral
+ validate_structured_output_request_outlines(params)
+ params.structured_outputs._backend = "outlines"
+ else:
+ # Fall back to guidance by default.
+ validate_guidance_grammar(params, tokenizer=None)
+ params.structured_outputs._backend = "guidance"
# Remember that this backend was set automatically
params.structured_outputs._backend_was_auto = True
diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py
index ad2ece50f9815..7e8ebe25c4603 100644
--- a/vllm/v1/executor/multiproc_executor.py
+++ b/vllm/v1/executor/multiproc_executor.py
@@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import (
get_dp_group,
get_ep_group,
get_inner_dp_world_group,
+ get_pcp_group,
get_pp_group,
get_tp_group,
)
@@ -110,12 +111,14 @@ class MultiprocExecutor(Executor):
f"({self.parallel_config.nnodes_within_dp}). "
)
self.local_world_size = self.parallel_config.local_world_size
- tensor_parallel_size = self.parallel_config.tensor_parallel_size
- pp_parallel_size = self.parallel_config.pipeline_parallel_size
- assert self.world_size == tensor_parallel_size * pp_parallel_size, (
+ tp_size = self.parallel_config.tensor_parallel_size
+ pp_size = self.parallel_config.pipeline_parallel_size
+ pcp_size = self.parallel_config.prefill_context_parallel_size
+ assert self.world_size == tp_size * pp_size * pcp_size, (
f"world_size ({self.world_size}) must be equal to the "
- f"tensor_parallel_size ({tensor_parallel_size}) x pipeline"
- f"_parallel_size ({pp_parallel_size}). "
+ f"tensor_parallel_size ({tp_size}) x pipeline"
+ f"_parallel_size ({pp_size}) x prefill_context"
+ f"_parallel_size ({pcp_size}). "
)
# Set multiprocessing envs
@@ -424,7 +427,11 @@ class MultiprocExecutor(Executor):
# 16-23, PP rank 2
# 24-31, PP rank 3
# so world_size - tp_size = 32 - 8 = 24 should be PP rank = -1 (i.e. 3)
- return self.world_size - self.parallel_config.tensor_parallel_size
+ return (
+ self.world_size
+ - self.parallel_config.tensor_parallel_size
+ * self.parallel_config.prefill_context_parallel_size
+ )
@dataclass
@@ -828,6 +835,8 @@ class WorkerProc:
dp_rank = get_dp_group().rank_in_group
pp_size = get_pp_group().world_size
pp_rank = get_pp_group().rank_in_group
+ pcp_size = get_pcp_group().world_size
+ pcp_rank = get_pcp_group().rank_in_group
tp_size = get_tp_group().world_size
tp_rank = get_tp_group().rank_in_group
dcp_size = get_dcp_group().world_size
@@ -837,6 +846,8 @@ class WorkerProc:
process_name += f"_DP{dp_rank}"
if pp_size > 1:
process_name += f"_PP{pp_rank}"
+ if pcp_size > 1:
+ process_name += f"_PCP{pcp_rank}"
if tp_size > 1:
process_name += f"_TP{tp_rank}"
if dcp_size > 1:
diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py
index 55db7445c9c74..406eafcd339b0 100644
--- a/vllm/v1/executor/ray_executor.py
+++ b/vllm/v1/executor/ray_executor.py
@@ -99,9 +99,9 @@ class RayDistributedExecutor(Executor):
# KV connector setup
self.has_connector = self.vllm_config.kv_transfer_config is not None
- self.ec_producer = (
- self.vllm_config.ec_transfer_config is not None
- and self.vllm_config.ec_transfer_config.is_ec_producer
+ self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
+ self.vllm_config.ec_transfer_config is None
+ or not self.vllm_config.ec_transfer_config.is_ec_producer
)
self.scheduler_output: SchedulerOutput | None = None
@@ -401,7 +401,7 @@ class RayDistributedExecutor(Executor):
"after execute_model() returns None."
)
- if self.ec_producer or not scheduler_output.total_num_scheduled_tokens:
+ if not self.uses_sampler or not scheduler_output.total_num_scheduled_tokens:
# Model will not execute, call model runner immediately.
return self._execute_dag(scheduler_output, None, non_block)
diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py
index 7f33eb7e699c7..751862aa9c767 100644
--- a/vllm/v1/kv_cache_interface.py
+++ b/vllm/v1/kv_cache_interface.py
@@ -95,10 +95,11 @@ class FullAttentionSpec(AttentionSpec):
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
max_model_len = vllm_config.model_config.max_model_len
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
+ pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
# Note(hc): each dcp rank only need save
# (max_model_len//dcp_world_size) tokens locally.
- if dcp_world_size > 1:
- max_model_len = cdiv(max_model_len, dcp_world_size)
+ if dcp_world_size * pcp_world_size > 1:
+ max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)
return cdiv(max_model_len, self.block_size) * self.page_size_bytes
@classmethod
diff --git a/vllm/v1/kv_offload/cpu.py b/vllm/v1/kv_offload/cpu.py
index 4b1bbe6f0cc2a..86747299eb107 100644
--- a/vllm/v1/kv_offload/cpu.py
+++ b/vllm/v1/kv_offload/cpu.py
@@ -4,8 +4,8 @@ from collections.abc import Iterator
import torch
-from vllm.config import VllmConfig, get_layers_from_vllm_config
-from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.attention import AttentionBackend
+from vllm.config import VllmConfig
from vllm.platforms import current_platform
from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
@@ -63,7 +63,9 @@ class CPUOffloadingSpec(OffloadingSpec):
return self._manager
def get_handlers(
- self, kv_caches: dict[str, torch.Tensor]
+ self,
+ kv_caches: dict[str, torch.Tensor],
+ attn_backends: dict[str, type[AttentionBackend]],
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
if not self._handler:
if not current_platform.is_cuda_alike():
@@ -71,15 +73,6 @@ class CPUOffloadingSpec(OffloadingSpec):
"CPU Offloading is currently only supported on CUDA-alike GPUs"
)
- layer_names = list(kv_caches.keys())
- layers = get_layers_from_vllm_config(
- self.vllm_config, AttentionLayerBase, layer_names
- )
- attn_backends = {
- layer_name: layers[layer_name].get_attn_backend()
- for layer_name in layer_names
- }
-
self._handler = CpuGpuOffloadingHandler(
attn_backends=attn_backends,
gpu_block_size=self.gpu_block_size,
diff --git a/vllm/v1/kv_offload/spec.py b/vllm/v1/kv_offload/spec.py
index a3c539a47d458..c1813a4ff4ea9 100644
--- a/vllm/v1/kv_offload/spec.py
+++ b/vllm/v1/kv_offload/spec.py
@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
from vllm.v1.kv_offload.worker.worker import OffloadingHandler
if TYPE_CHECKING:
+ from vllm.attention import AttentionBackend
from vllm.config import VllmConfig
logger = init_logger(__name__)
@@ -48,13 +49,16 @@ class OffloadingSpec(ABC):
@abstractmethod
def get_handlers(
- self, kv_caches: dict[str, torch.Tensor]
+ self,
+ kv_caches: dict[str, torch.Tensor],
+ attn_backends: dict[str, type["AttentionBackend"]],
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
"""
Get offloading handlers along with their respective src and dst types.
Args:
kv_caches: A dictionary of layer_name -> gpu_kv_cache tensor.
+ attn_backends: A dictionary of layer_name -> AttentionBackend.
Yields:
Tuples of (src_type, dst_type, offloading_handler).
diff --git a/vllm/v1/kv_offload/worker/cpu_gpu.py b/vllm/v1/kv_offload/worker/cpu_gpu.py
index 646f9d0d75423..bb163f0043fc6 100644
--- a/vllm/v1/kv_offload/worker/cpu_gpu.py
+++ b/vllm/v1/kv_offload/worker/cpu_gpu.py
@@ -68,9 +68,9 @@ class CpuGpuOffloadingHandler(OffloadingHandler):
self.h2d_stream = torch.cuda.Stream()
# job_id -> transfer cuda event
- self.transfer_events: dict[int, torch.cuda.Event] = {}
+ self.transfer_events: dict[int, torch.Event] = {}
# list of cuda events available for re-use
- self.events_pool: list[torch.cuda.Event] = []
+ self.events_pool: list[torch.Event] = []
pin_memory = is_pin_memory_available()
@@ -83,10 +83,18 @@ class CpuGpuOffloadingHandler(OffloadingHandler):
self.gpu_tensors.append(gpu_tensor)
gpu_shape = gpu_tensor.shape
- test_shape = attn_backends[layer_name].get_kv_cache_shape(
+ attn_backend = attn_backends[layer_name]
+ test_shape = attn_backend.get_kv_cache_shape(
num_blocks=1234, block_size=16, num_kv_heads=8, head_size=256
)
- if test_shape[0] == 1234:
+
+ if len(gpu_shape) != len(test_shape):
+ # cross-layers tensor
+ # shape is (num_blocks, ...)
+ assert len(gpu_shape) == len(test_shape) + 1
+ num_blocks_idx = 0
+ self.kv_dim_before_num_blocks.append(False)
+ elif test_shape[0] == 1234:
# shape is (num_blocks, ...)
num_blocks_idx = 0
self.kv_dim_before_num_blocks.append(False)
@@ -135,25 +143,23 @@ class CpuGpuOffloadingHandler(OffloadingHandler):
assert src_blocks.ndim == 1
assert dst_blocks.ndim == 1
- dst_sub_blocks_to_skip = -src_blocks.size % dst_block_size_factor
src_sub_block_count = src_blocks.size * src_block_size_factor
+ dst_sub_block_count = dst_blocks.size * dst_block_size_factor
+ src_sub_blocks_to_skip = -dst_blocks.size % src_block_size_factor
- assert (
- src_sub_block_count
- == dst_blocks.size * dst_block_size_factor - dst_sub_blocks_to_skip
- )
+ assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip
- src_to_dst = np.empty((src_sub_block_count, 2), dtype=np.int64)
- expand_block_ids(src_blocks, src_block_size_factor, src_to_dst[:, 0])
+ src_to_dst = np.empty((dst_sub_block_count, 2), dtype=np.int64)
expand_block_ids(
- dst_blocks,
- dst_block_size_factor,
- src_to_dst[:, 1],
- skip_count=dst_sub_blocks_to_skip,
+ src_blocks,
+ src_block_size_factor,
+ src_to_dst[:, 0],
+ skip_count=src_sub_blocks_to_skip,
)
+ expand_block_ids(dst_blocks, dst_block_size_factor, src_to_dst[:, 1])
src_to_dst_tensor = torch.from_numpy(src_to_dst)
- event = self.events_pool.pop() if self.events_pool else torch.cuda.Event()
+ event = self.events_pool.pop() if self.events_pool else torch.Event()
with torch.cuda.stream(stream):
for src_tensor, dst_tensor, kv_dim in zip(
src_tensors, dst_tensors, self.kv_dim_before_num_blocks
diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py
index cb36e7973650e..e2d82241ce210 100644
--- a/vllm/v1/metrics/loggers.py
+++ b/vllm/v1/metrics/loggers.py
@@ -16,7 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
KVConnectorPrometheus,
)
from vllm.logger import init_logger
-from vllm.plugins import load_plugins_by_group
+from vllm.plugins import STAT_LOGGER_PLUGINS_GROUP, load_plugins_by_group
from vllm.v1.engine import FinishReason
from vllm.v1.metrics.prometheus import unregister_vllm_metrics
from vllm.v1.metrics.stats import (
@@ -67,7 +67,7 @@ class StatLoggerBase(ABC):
def load_stat_logger_plugin_factories() -> list[StatLoggerFactory]:
factories: list[StatLoggerFactory] = []
- for name, plugin_class in load_plugins_by_group("vllm.stat_logger_plugins").items():
+ for name, plugin_class in load_plugins_by_group(STAT_LOGGER_PLUGINS_GROUP).items():
if not isinstance(plugin_class, type) or not issubclass(
plugin_class, StatLoggerBase
):
diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py
index c0b2835c3124c..e32d5bb608b1d 100644
--- a/vllm/v1/outputs.py
+++ b/vllm/v1/outputs.py
@@ -158,7 +158,7 @@ class ModelRunnerOutput:
# num_generated_tokens is the number of tokens
# generated in the current step. It can be different for
# each request due to speculative/jump decoding.
- sampled_token_ids: list[np.ndarray]
+ sampled_token_ids: list[list[int]]
# [num_reqs, max_num_logprobs + 1]
# [num_reqs, max_num_logprobs + 1]
@@ -220,7 +220,7 @@ def make_empty_encoder_model_runner_output(
req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)}
# No tokens generated yet ⇒ one empty list per request
- sampled_token_ids: list[list[int]] = [np.array([0]) for _ in req_ids]
+ sampled_token_ids: list[list[int]] = [[0] for _ in req_ids]
# Pooler outputs are not available yet ⇒ use None placeholders
pooler_output: list[torch.Tensor | None] = [None for _ in req_ids]
diff --git a/vllm/v1/request.py b/vllm/v1/request.py
index 3d92906fbf4b1..366cdadf5a583 100644
--- a/vllm/v1/request.py
+++ b/vllm/v1/request.py
@@ -121,6 +121,9 @@ class Request:
# The number of requests being preempted by the scheduler
self.num_preemptions = 0
+ # The number of tokens that have been computed remotely.
+ self.num_external_computed_tokens = 0
+
self.block_hashes: list[BlockHash] = []
self.get_hash_new_full_blocks: Callable[[], list[BlockHash]] | None = None
if block_hasher is not None:
diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py
index 02ea658b7f20e..5b2d130b0ea42 100644
--- a/vllm/v1/sample/ops/topk_topp_sampler.py
+++ b/vllm/v1/sample/ops/topk_topp_sampler.py
@@ -7,6 +7,7 @@ import torch.nn as nn
from packaging import version
from vllm import envs
+from vllm._aiter_ops import rocm_aiter_ops
from vllm.config.model import LogprobsMode
from vllm.logger import init_logger
from vllm.platforms import CpuArchEnum, current_platform
@@ -55,6 +56,24 @@ class TopKTopPSampler(nn.Module):
self.forward = self.forward_native
else:
self.forward = self.forward_cpu
+ elif (
+ logprobs_mode not in ("processed_logits", "processed_logprobs")
+ and rocm_aiter_ops.is_enabled()
+ ):
+ try:
+ import aiter.ops.sampling # noqa: F401
+
+ self.aiter_ops = torch.ops.aiter
+ logger.info_once(
+ "Using aiter sampler on ROCm (lazy import, sampling-only)."
+ )
+ self.forward = self.forward_hip
+ except ImportError:
+ logger.warning_once(
+ "aiter.ops.sampling is not available on ROCm. "
+ "Falling back to forward_native implementation."
+ )
+ self.forward = self.forward_native
else:
self.forward = self.forward_native
@@ -138,6 +157,64 @@ class TopKTopPSampler(nn.Module):
return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return
+ def forward_hip(
+ self,
+ logits: torch.Tensor,
+ generators: dict[int, torch.Generator],
+ k: torch.Tensor | None,
+ p: torch.Tensor | None,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Optimized ROCm/aiter path (same structure as forward_cuda)."""
+ if (k is None and p is None) or generators:
+ if generators:
+ logger.warning_once(
+ "aiter sampler does not support per-request generators; "
+ "falling back to PyTorch-native."
+ )
+ return self.forward_native(logits, generators, k, p)
+ assert self.logprobs_mode not in (
+ "processed_logits",
+ "processed_logprobs",
+ ), "aiter sampler does not support returning logits/logprobs."
+ return self.aiter_sample(logits, k, p, generators), None
+
+ def aiter_sample(
+ self,
+ logits: torch.Tensor,
+ k: torch.Tensor | None,
+ p: torch.Tensor | None,
+ generators: dict[int, torch.Generator],
+ ) -> torch.Tensor:
+ """Sample from logits using aiter ops."""
+ use_top_k = k is not None
+ use_top_p = p is not None
+ # Joint k+p path
+ if use_top_p and use_top_k:
+ probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous()
+ next_token_ids = self.aiter_ops.top_k_top_p_sampling_from_probs(
+ probs,
+ None,
+ *_to_tensor_scalar_tuple(k),
+ *_to_tensor_scalar_tuple(p),
+ deterministic=True,
+ )
+ return next_token_ids.view(-1)
+ # Top-p only path
+ elif use_top_p:
+ probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous()
+ next_token_ids = self.aiter_ops.top_p_sampling_from_probs(
+ probs, None, *_to_tensor_scalar_tuple(p), deterministic=True
+ )
+ return next_token_ids.view(-1)
+ # Top-k only path
+ elif use_top_k:
+ probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous()
+ renorm_probs = self.aiter_ops.top_k_renorm_probs(
+ probs, *_to_tensor_scalar_tuple(k)
+ )
+ return torch.multinomial(renorm_probs, num_samples=1).view(-1)
+ raise RuntimeError("aiter_sample was called with no active top-k or top-p.")
+
# Note: this is a workaround for
# https://github.com/pytorch/pytorch/pull/151218
@@ -288,3 +365,10 @@ def flashinfer_sample(
)
return next_token_ids.view(-1)
+
+
+def _to_tensor_scalar_tuple(x):
+ if isinstance(x, torch.Tensor):
+ return (x, 0)
+ else:
+ return (None, x)
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index f31a0cddda9ae..926305d25f56b 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/vllm/v1/sample/rejection_sampler.py
@@ -3,7 +3,6 @@
from dataclasses import replace
-import numpy as np
import torch
import torch.nn as nn
@@ -205,7 +204,7 @@ class RejectionSampler(nn.Module):
def parse_output(
output_token_ids: torch.Tensor,
vocab_size: int,
- ) -> list[np.ndarray]:
+ ) -> list[list[int]]:
"""Parse the output of the rejection sampler.
Args:
output_token_ids: The sampled token IDs in shape
@@ -221,7 +220,10 @@ class RejectionSampler(nn.Module):
valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & (
output_token_ids_np < vocab_size
)
- return [row[valid_mask[i]] for i, row in enumerate(output_token_ids_np)]
+ outputs = [
+ row[valid_mask[i]].tolist() for i, row in enumerate(output_token_ids_np)
+ ]
+ return outputs
def apply_logits_processors(
self,
diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py
index 39c63fe31ad2c..c75b4f0543c0d 100644
--- a/vllm/v1/sample/sampler.py
+++ b/vllm/v1/sample/sampler.py
@@ -81,7 +81,10 @@ class Sampler(nn.Module):
if logprobs_mode == "raw_logprobs":
raw_logprobs = self.compute_logprobs(logits)
elif logprobs_mode == "raw_logits":
- raw_logprobs = logits.clone()
+ if logits.dtype == torch.float32:
+ raw_logprobs = logits.clone()
+ else:
+ raw_logprobs = logits.to(torch.float32)
# Use float32 for the logits.
logits = logits.to(torch.float32)
diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py
index 5bf2503c3027d..3de418f1d13c8 100644
--- a/vllm/v1/spec_decode/eagle.py
+++ b/vllm/v1/spec_decode/eagle.py
@@ -86,9 +86,9 @@ class EagleProposer:
self.use_cuda_graph = False
- compilation_config = self.vllm_config.compilation_config
- if compilation_config.mode == CompilationMode.VLLM_COMPILE:
- cudagraph_mode = compilation_config.cudagraph_mode
+ self.compilation_config = self.vllm_config.compilation_config
+ if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:
+ cudagraph_mode = self.compilation_config.cudagraph_mode
if cudagraph_mode != CUDAGraphMode.NONE and not cudagraph_mode.has_mode(
CUDAGraphMode.PIECEWISE
):
@@ -103,22 +103,24 @@ class EagleProposer:
and not self.speculative_config.enforce_eager
)
- self.cudagraph_batch_sizes = (
- (sorted(self.vllm_config.compilation_config.cudagraph_capture_sizes))
- if self.use_cuda_graph
- else []
- )
-
- self.use_cuda_graph = self.use_cuda_graph and bool(self.cudagraph_batch_sizes)
# persistent buffers for cuda graph
self.input_ids = torch.zeros(
self.max_num_tokens, dtype=torch.int32, device=device
)
self.uses_mrope = self.vllm_config.model_config.uses_mrope
if self.uses_mrope:
- # M-RoPE need (3, max_num_tokens)
+ # NOTE: `mrope_positions` is implemented with one additional dummy
+ # position on purpose to make it non-contiguous so that it can work
+ # with torch compile.
+ # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
+
+ # NOTE: When M-RoPE is enabled, position ids are 3D regardless of
+ # the modality of inputs. For text-only inputs, each dimension has
+ # identical position IDs, making M-RoPE functionally equivalent to
+ # 1D-RoPE.
+ # See page 5 of https://arxiv.org/abs/2409.12191
self.mrope_positions = torch.zeros(
- (3, self.max_num_tokens), dtype=torch.int64, device=device
+ (3, self.max_num_tokens + 1), dtype=torch.int64, device=device
)
else:
# RoPE need (max_num_tokens,)
@@ -267,7 +269,10 @@ class EagleProposer:
per_layer_attn_metadata[layer_name] = draft_indexer_metadata
cudagraph_runtime_mode = CUDAGraphMode.NONE
- if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
+ if (
+ self.use_cuda_graph
+ and num_tokens <= self.compilation_config.max_cudagraph_capture_size
+ ):
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
else:
@@ -357,7 +362,10 @@ class EagleProposer:
# Generate the remaining draft tokens.
draft_token_ids_list = [draft_token_ids]
- if self.use_cuda_graph and batch_size <= self.cudagraph_batch_sizes[-1]:
+ if (
+ self.use_cuda_graph
+ and batch_size <= self.compilation_config.max_cudagraph_capture_size
+ ):
input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
else:
@@ -487,7 +495,7 @@ class EagleProposer:
def prepare_next_token_ids_cpu(
self,
- sampled_token_ids: list[np.ndarray],
+ sampled_token_ids: list[list[int]],
requests: dict[str, CachedRequestState],
gpu_input_batch: InputBatch,
num_scheduled_tokens: dict[str, int],
@@ -502,7 +510,7 @@ class EagleProposer:
req_ids = gpu_input_batch.req_ids
next_token_ids: list[int] = []
for i, token_ids in enumerate(sampled_token_ids):
- if token_ids.shape[0] > 0:
+ if token_ids:
# Common case.
next_token_id = token_ids[-1]
else:
@@ -513,9 +521,10 @@ class EagleProposer:
seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]
next_token_id = req_state.get_token_id(seq_len)
next_token_ids.append(next_token_id)
- return torch.tensor(
+ next_token_ids = torch.tensor(
next_token_ids, dtype=torch.int32, device=self.input_ids.device
)
+ return next_token_ids
def prepare_next_token_ids_padded(
self,
@@ -767,7 +776,10 @@ class EagleProposer:
self.positions[:num_tokens] = tree_positions.view(-1)
self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)
- if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
+ if (
+ self.use_cuda_graph
+ and num_tokens <= self.compilation_config.max_cudagraph_capture_size
+ ):
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
else:
@@ -1019,8 +1031,11 @@ class EagleProposer:
elif (
isinstance(target_embed_tokens.weight, torch.Tensor)
and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)
- and torch.equal(
- target_embed_tokens.weight, self.model.model.embed_tokens.weight
+ and torch.allclose(
+ target_embed_tokens.weight.cpu(),
+ self.model.model.embed_tokens.weight.cpu(),
+ rtol=1e-5,
+ atol=1e-7,
)
):
share_embeddings = True
@@ -1101,7 +1116,10 @@ class EagleProposer:
) -> None:
# Determine if CUDA graphs should be used for this run.
cudagraphs_enabled = use_cudagraphs and self.use_cuda_graph
- if cudagraphs_enabled and num_tokens <= self.cudagraph_batch_sizes[-1]:
+ if (
+ cudagraphs_enabled
+ and num_tokens <= self.compilation_config.max_cudagraph_capture_size
+ ):
num_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
with set_forward_context(
diff --git a/vllm/v1/spec_decode/ngram_proposer.py b/vllm/v1/spec_decode/ngram_proposer.py
index 378937dba9882..e2f83cb24aa90 100644
--- a/vllm/v1/spec_decode/ngram_proposer.py
+++ b/vllm/v1/spec_decode/ngram_proposer.py
@@ -54,7 +54,7 @@ class NgramProposer:
# Trigger Numba JIT compilation for N-gram proposer.
# This usually takes less than 1 second.
self.propose(
- [np.array([])] * 1024,
+ [[]] * 1024,
[""] * 1024,
np.zeros(1024, dtype=np.int32),
np.zeros((1024, self.max_model_len), dtype=np.int32),
@@ -131,7 +131,7 @@ class NgramProposer:
def propose(
self,
- sampled_token_ids: list[np.ndarray],
+ sampled_token_ids: list[list[int]],
req_ids: list[str],
num_tokens_no_spec: np.ndarray,
token_ids_cpu: np.ndarray,
@@ -140,7 +140,7 @@ class NgramProposer:
# find which requests need ngram proposals
valid_ngram_requests = []
for i, sampled_ids in enumerate(sampled_token_ids):
- num_sampled_ids = sampled_ids.shape[0]
+ num_sampled_ids = len(sampled_ids)
if not num_sampled_ids:
# Skip speculative decoding.
continue
diff --git a/vllm/v1/spec_decode/suffix_decoding.py b/vllm/v1/spec_decode/suffix_decoding.py
index d76e0ffe778d4..049e335db3254 100644
--- a/vllm/v1/spec_decode/suffix_decoding.py
+++ b/vllm/v1/spec_decode/suffix_decoding.py
@@ -1,7 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import numpy as np
-
from vllm.config import VllmConfig
from vllm.v1.worker.gpu_input_batch import InputBatch
@@ -34,16 +32,16 @@ class SuffixDecodingProposer:
def propose(
self,
input_batch: InputBatch,
- sampled_token_ids: list[np.ndarray],
+ sampled_token_ids: list[list[int]],
) -> list[list[int]]:
"""
Propose speculative tokens for each request in the input batch. Suffix Decoding
will speculate a dynamic number of tokens for each request every decoding step,
so each entry in the returned list may have different lengths.
"""
- draft_token_ids: list[np.ndarray] = []
+ draft_token_ids: list[list[int]] = []
for i, sampled_ids in enumerate(sampled_token_ids):
- if sampled_ids.shape[0] == 0:
+ if not sampled_ids:
# Skip speculative decoding for partial prefills.
draft_token_ids.append([])
continue
@@ -72,7 +70,7 @@ class SuffixDecodingProposer:
self.suffix_cache.start_request(req_id, prompt_token_ids)
# Append the newly sampled ids to the suffix cache for this request.
- self.suffix_cache.add_active_response(req_id, sampled_ids.tolist())
+ self.suffix_cache.add_active_response(req_id, sampled_ids)
# Suffix decoding only uses the most recent tokens up to max_tree_depth, so
# we extract the pattern from the end of the input.
diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py
index 9f6c19e464308..37ec0fb97e06b 100644
--- a/vllm/v1/worker/block_table.py
+++ b/vllm/v1/worker/block_table.py
@@ -4,7 +4,7 @@
import numpy as np
import torch
-from vllm.distributed import get_dcp_group
+from vllm.distributed import get_dcp_group, get_pcp_group
from vllm.logger import init_logger
from vllm.utils.math_utils import cdiv
from vllm.v1.utils import CpuGpuBuffer
@@ -22,7 +22,7 @@ class BlockTable:
pin_memory: bool,
device: torch.device,
kernel_block_size: int,
- dcp_kv_cache_interleave_size: int,
+ cp_kv_cache_interleave_size: int,
):
"""
Args:
@@ -80,6 +80,13 @@ class BlockTable:
else:
self._kernel_block_arange = None
+ try:
+ self.pcp_world_size = get_pcp_group().world_size
+ self.pcp_rank = get_pcp_group().rank_in_group
+ except AssertionError:
+ # PCP might not be initialized in testing
+ self.pcp_world_size = 1
+ self.pcp_rank = 0
try:
self.dcp_world_size = get_dcp_group().world_size
self.dcp_rank = get_dcp_group().rank_in_group
@@ -87,7 +94,7 @@ class BlockTable:
# DCP might not be initialized in testing
self.dcp_world_size = 1
self.dcp_rank = 0
- self.dcp_kv_cache_interleave_size = dcp_kv_cache_interleave_size
+ self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
def append_row(
self,
@@ -131,14 +138,16 @@ class BlockTable:
# NOTE(woosuk): We can't simply use `token_indices // block_size`
# here because M (max_model_len) is not necessarily divisible by
# block_size.
- if self.dcp_world_size > 1:
+ total_cp_world_size = self.pcp_world_size * self.dcp_world_size
+ total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
+ if total_cp_world_size > 1:
# Note(hc): The DCP implement store kvcache with an interleave
# style, the kvcache for the token whose token_idx is i is
# always stored on the GPU whose dcp_rank equals i % cp_world_size:
# Use a "virtual block" which equals to world_size * block_size
# for block_table_indices calculation.
- virtual_block_size = self.block_size * self.dcp_world_size
+ virtual_block_size = self.block_size * total_cp_world_size
block_table_indices = (
req_indices * self.max_num_blocks_per_req
+ positions // virtual_block_size
@@ -150,16 +159,16 @@ class BlockTable:
virtual_block_offsets = positions % virtual_block_size
mask = (
virtual_block_offsets
- // self.dcp_kv_cache_interleave_size
- % self.dcp_world_size
- == self.dcp_rank
+ // self.cp_kv_cache_interleave_size
+ % total_cp_world_size
+ == total_cp_rank
)
# Calculate local block_offsets
block_offsets = (
virtual_block_offsets
- // (self.dcp_world_size * self.dcp_kv_cache_interleave_size)
- * self.dcp_kv_cache_interleave_size
- + virtual_block_offsets % self.dcp_kv_cache_interleave_size
+ // (total_cp_world_size * self.cp_kv_cache_interleave_size)
+ * self.cp_kv_cache_interleave_size
+ + virtual_block_offsets % self.cp_kv_cache_interleave_size
)
# Calculate slot_mapping
slot_mapping = block_numbers * self.block_size + block_offsets
@@ -253,12 +262,17 @@ class MultiGroupBlockTable:
block_sizes: list[int],
kernel_block_sizes: list[int],
num_speculative_tokens: int = 0,
- dcp_kv_cache_interleave_size: int = 1,
+ cp_kv_cache_interleave_size: int = 1,
) -> None:
# Note(hc): each dcp rank only store
# (max_model_len//dcp_world_size) tokens in kvcache,
# so the block_size which used for calc max_num_blocks_per_req
# must be multiplied by dcp_world_size.
+ try:
+ pcp_world_size = get_pcp_group().world_size
+ except AssertionError:
+ # PCP might not be initialized in testing
+ pcp_world_size = 1
try:
dcp_world_size = get_dcp_group().world_size
except AssertionError:
@@ -271,19 +285,21 @@ class MultiGroupBlockTable:
f"must match block_sizes length ({len(block_sizes)})"
)
+ total_cp_world_size = dcp_world_size * pcp_world_size
+
self.block_tables = [
BlockTable(
block_size,
max_num_reqs,
max(
- cdiv(max_model_len, block_size * dcp_world_size),
+ cdiv(max_model_len, block_size * total_cp_world_size),
1 + num_speculative_tokens,
),
max_num_batched_tokens,
pin_memory,
device,
kernel_block_size,
- dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size,
)
for block_size, kernel_block_size in zip(block_sizes, kernel_block_sizes)
]
diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py
index 40f011fed1ada..6bfbc32d598fa 100644
--- a/vllm/v1/worker/cpu_model_runner.py
+++ b/vllm/v1/worker/cpu_model_runner.py
@@ -96,14 +96,14 @@ def _torch_cuda_wrapper():
def __init__(self, *args, **kwargs) -> None:
pass
- cuda_event = torch.cuda.Event
+ cuda_event = torch.Event
cuda_stream = torch.cuda.Stream
try:
- torch.cuda.Event = _EventPlaceholder
+ torch.Event = _EventPlaceholder
torch.cuda.Stream = _StreamPlaceholder
yield
finally:
- torch.cuda.Event = cuda_event
+ torch.Event = cuda_event
torch.cuda.Stream = cuda_stream
diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py
index 4420a057d1e58..b080fea1d2dd6 100644
--- a/vllm/v1/worker/cpu_worker.py
+++ b/vllm/v1/worker/cpu_worker.py
@@ -3,6 +3,7 @@
import os
import platform
from collections.abc import Callable
+from typing import Any
import torch
@@ -37,6 +38,9 @@ class CPUWorker(Worker):
self.parallel_config.disable_custom_all_reduce = True
+ # Torch profiler. Enabled and configured through env vars:
+ # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+ self.profiler: Any | None = None
if envs.VLLM_TORCH_PROFILER_DIR:
torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
@@ -80,13 +84,13 @@ class CPUWorker(Worker):
self.local_omp_cpuid = "nobind"
else:
local_dp_rank = self.parallel_config.data_parallel_rank_local
- omp_cpuids = omp_cpuids.split("|")
+ omp_cpuids_list = omp_cpuids.split("|")
if local_dp_rank is not None:
world_size = self.parallel_config.world_size
- omp_cpuids = omp_cpuids[
+ omp_cpuids_list = omp_cpuids_list[
local_dp_rank * world_size : (local_dp_rank + 1) * world_size
]
- self.local_omp_cpuid = omp_cpuids[self.rank]
+ self.local_omp_cpuid = omp_cpuids_list[self.rank]
if self.local_omp_cpuid != "nobind":
ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
@@ -120,7 +124,7 @@ class CPUWorker(Worker):
pass
def determine_available_memory(self) -> int:
- return self.cache_config.cpu_kvcache_space_bytes # type: ignore
+ return self.cache_config.cpu_kvcache_space_bytes or 0
def compile_or_warm_up_model(self) -> None:
# Reset the seed to ensure that the random state is not affected by
diff --git a/vllm/v1/worker/gpu/README.md b/vllm/v1/worker/gpu/README.md
new file mode 100644
index 0000000000000..093f524b3250f
--- /dev/null
+++ b/vllm/v1/worker/gpu/README.md
@@ -0,0 +1,4 @@
+# [Experimental] Model Runner V2
+
+This directory contains the new model runner which is under active development.
+Ping [Woosuk Kwon](https://github.com/WoosukKwon) for any changes.
diff --git a/vllm/v1/worker/gpu/__init__.py b/vllm/v1/worker/gpu/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py
new file mode 100644
index 0000000000000..421fb29a7f87f
--- /dev/null
+++ b/vllm/v1/worker/gpu/async_utils.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from contextlib import contextmanager
+
+import torch
+
+from vllm.v1.outputs import (
+ AsyncModelRunnerOutput,
+ LogprobsTensors,
+ ModelRunnerOutput,
+ SamplerOutput,
+)
+
+
+class AsyncOutput(AsyncModelRunnerOutput):
+ def __init__(
+ self,
+ model_runner_output: ModelRunnerOutput,
+ sampler_output: SamplerOutput,
+ num_sampled_tokens: torch.Tensor,
+ copy_stream: torch.cuda.Stream,
+ copy_event: torch.cuda.Event,
+ ):
+ self.model_runner_output = model_runner_output
+ self.sampler_output = sampler_output
+ self.num_sampled_tokens = num_sampled_tokens
+ self.copy_stream = copy_stream
+ self.copy_event = copy_event
+
+ default_stream = torch.cuda.current_stream()
+ with torch.cuda.stream(self.copy_stream):
+ self.copy_stream.wait_stream(default_stream)
+
+ # NOTE(woosuk): We must ensure that CPU tensors are not freed
+ # before the device-to-host copy is fully completed. For instance,
+ # operations like
+ # self.sampled_token_np = ...to("cpu", non_blocking=True).numpy()
+ # are unsafe because the underlying CPU tensor can be prematurely freed and
+ # reused by other tensors before the asynchronous copy finishes, potentially
+ # causing race conditions. To prevent this, we delay freeing by holding
+ # references until the copy event signals completion.
+ # Likewise, we also need to keep the reference to the GPU tensors.
+ # This is done by keeping the reference to sampler_output and
+ # model_runner_output.
+ self.sampled_token_ids = sampler_output.sampled_token_ids.to(
+ "cpu", non_blocking=True
+ )
+ if sampler_output.logprobs_tensors is not None:
+ self.logprobs_tensors: LogprobsTensors | None = (
+ sampler_output.logprobs_tensors.to_cpu_nonblocking()
+ )
+ else:
+ self.logprobs_tensors = None
+ self.num_sampled_tokens = num_sampled_tokens.to("cpu", non_blocking=True)
+ self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
+ if self.model_runner_output.prompt_logprobs_dict:
+ for k, v in self.model_runner_output.prompt_logprobs_dict.items():
+ if v is not None:
+ self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
+ else:
+ self.prompt_logprobs_dict[k] = None
+ self.copy_event.record(self.copy_stream)
+
+ def get_output(self) -> ModelRunnerOutput:
+ self.copy_event.synchronize()
+ num_sampled_tokens_np = self.num_sampled_tokens.numpy()
+
+ # NOTE(woosuk): The following code is to ensure compatibility with
+ # the existing model runner.
+ # Going forward, we should keep the data structures as NumPy arrays
+ # rather than Python lists.
+ sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
+ num_reqs = len(sampled_token_ids)
+ for i in range(num_reqs):
+ del sampled_token_ids[i][num_sampled_tokens_np[i] :]
+ self.model_runner_output.sampled_token_ids = sampled_token_ids
+
+ if self.logprobs_tensors is not None:
+ self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
+ self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
+ return self.model_runner_output
+
+
+@contextmanager
+def async_barrier(event: torch.cuda.Event | None):
+ if event is not None:
+ event.synchronize()
+ try:
+ yield
+ finally:
+ if event is not None:
+ event.record()
diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py
new file mode 100644
index 0000000000000..4510a1c5ca1e9
--- /dev/null
+++ b/vllm/v1/worker/gpu/attn_utils.py
@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Sequence
+from typing import Any, cast
+
+import numpy as np
+import torch
+
+from vllm.attention.backends.abstract import AttentionBackend
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.v1.attention.backends.utils import (
+ AttentionMetadataBuilder,
+ CommonAttentionMetadata,
+)
+from vllm.v1.kv_cache_interface import (
+ AttentionSpec,
+ KVCacheConfig,
+ KVCacheSpec,
+)
+from vllm.v1.utils import CpuGpuBuffer
+from vllm.v1.worker.utils import bind_kv_cache
+
+
+def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
+ kv_cache_spec: dict[str, KVCacheSpec] = {}
+ layer_type = cast(type[Any], AttentionLayerBase)
+ attn_layers = get_layers_from_vllm_config(vllm_config, layer_type)
+ for layer_name, attn_module in attn_layers.items():
+ # Skip modules that don't need KV cache (eg encoder-only attention)
+ if spec := attn_module.get_kv_cache_spec(vllm_config):
+ kv_cache_spec[layer_name] = spec
+ return kv_cache_spec
+
+
+def init_attn_backend(
+ kv_cache_config: KVCacheConfig,
+ vllm_config: VllmConfig,
+ device: torch.device,
+):
+ attn_backends: dict[str, type[AttentionBackend]] = {}
+ attn_metadata_builders: list[AttentionMetadataBuilder] = []
+ flashinfer_workspace: torch.Tensor | None = None
+ for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
+ layer_names = kv_cache_group_spec.layer_names
+ any_layer_name = next(iter(layer_names))
+
+ layer_type = cast(type[Any], AttentionLayerBase)
+ attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
+ attn_backend = attn_layers[any_layer_name].get_attn_backend()
+ for layer_name in layer_names:
+ attn_backends[layer_name] = attn_backend
+
+ attn_metadata_builder = attn_backend.get_builder_cls()(
+ kv_cache_group_spec.kv_cache_spec,
+ layer_names,
+ vllm_config,
+ device,
+ )
+ attn_metadata_builders.append(attn_metadata_builder) # type: ignore
+
+ if "FLASHINFER" in attn_backend.get_name():
+ if flashinfer_workspace is None:
+ flashinfer_workspace = attn_metadata_builder._get_workspace_buffer()
+ else:
+ attn_metadata_builder.set_workspace_buffer(flashinfer_workspace)
+ return attn_backends, attn_metadata_builders
+
+
+def _allocate_kv_cache(
+ kv_cache_config: KVCacheConfig,
+ device: torch.device,
+):
+ kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
+ for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
+ tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
+ for layer_name in kv_cache_tensor.shared_by:
+ kv_cache_raw_tensors[layer_name] = tensor
+
+ layer_names = set()
+ for group in kv_cache_config.kv_cache_groups:
+ for layer_name in group.layer_names:
+ layer_names.add(layer_name)
+ assert layer_names == set(kv_cache_raw_tensors.keys()), (
+ "Some layers are not correctly initialized"
+ )
+ return kv_cache_raw_tensors
+
+
+def _reshape_kv_cache(
+ kv_cache_config: KVCacheConfig,
+ kv_cache_raw_tensors: dict[str, torch.Tensor],
+ attn_backends: dict[str, AttentionBackend],
+) -> dict[str, torch.Tensor]:
+ kv_caches: dict[str, torch.Tensor] = {}
+ for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
+ kv_cache_spec = kv_cache_group_spec.kv_cache_spec
+ assert isinstance(kv_cache_spec, AttentionSpec)
+ for layer_name in kv_cache_group_spec.layer_names:
+ raw_tensor = kv_cache_raw_tensors[layer_name]
+ assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
+ num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
+
+ attn_backend = attn_backends[layer_name]
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
+ num_blocks,
+ kv_cache_spec.block_size,
+ kv_cache_spec.num_kv_heads,
+ kv_cache_spec.head_size,
+ )
+
+ # FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
+ try:
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
+ assert len(kv_cache_stride_order) == len(kv_cache_shape)
+ except (AttributeError, NotImplementedError):
+ kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+ kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+ inv_order = [
+ kv_cache_stride_order.index(i)
+ for i in range(len(kv_cache_stride_order))
+ ]
+
+ dtype = kv_cache_spec.dtype
+ raw_tensor = raw_tensor.view(dtype)
+ raw_tensor = raw_tensor.view(kv_cache_shape)
+ kv_caches[layer_name] = raw_tensor.permute(*inv_order)
+ return kv_caches
+
+
+def init_kv_cache(
+ runner_kv_caches: list[torch.Tensor],
+ forward_context: dict[str, Any],
+ kv_cache_config: KVCacheConfig,
+ attn_backends: dict[str, AttentionBackend],
+ device: torch.device,
+) -> None:
+ kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
+ kv_caches = _reshape_kv_cache(kv_cache_config, kv_cache_raw_tensors, attn_backends)
+ bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
+
+
+def build_attn_metadata(
+ attn_metadata_builders: list[AttentionMetadataBuilder],
+ num_reqs: int,
+ num_tokens: int,
+ query_start_loc: CpuGpuBuffer,
+ seq_lens: torch.Tensor,
+ seq_lens_np: np.ndarray,
+ num_computed_tokens_cpu: torch.Tensor | None,
+ block_tables: Sequence[torch.Tensor],
+ slot_mappings: torch.Tensor,
+ kv_cache_config: KVCacheConfig,
+) -> dict[str, Any]:
+ query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
+ query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
+ max_query_len = int(query_start_loc.np[: num_reqs + 1].max())
+ seq_lens = seq_lens[:num_reqs]
+ seq_lens_cpu = torch.from_numpy(seq_lens_np)
+ max_seq_len = int(seq_lens_np.max())
+
+ attn_metadata: dict[str, Any] = {}
+ kv_cache_groups = kv_cache_config.kv_cache_groups
+ for i, kv_cache_spec in enumerate(kv_cache_groups):
+ block_table = block_tables[i]
+ slot_mapping = slot_mappings[i]
+
+ common_attn_metadata = CommonAttentionMetadata(
+ query_start_loc=query_start_loc_gpu,
+ query_start_loc_cpu=query_start_loc_cpu,
+ seq_lens=seq_lens,
+ seq_lens_cpu=seq_lens_cpu,
+ max_seq_len=max_seq_len,
+ num_computed_tokens_cpu=num_computed_tokens_cpu,
+ num_reqs=num_reqs,
+ num_actual_tokens=num_tokens,
+ max_query_len=max_query_len,
+ block_table_tensor=block_table,
+ slot_mapping=slot_mapping,
+ causal=True,
+ )
+
+ attn_metadata_builder = attn_metadata_builders[i]
+ metadata = attn_metadata_builder.build(
+ common_prefix_len=0,
+ common_attn_metadata=common_attn_metadata,
+ )
+ for layer_name in kv_cache_spec.layer_names:
+ attn_metadata[layer_name] = metadata
+ return attn_metadata
diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py
new file mode 100644
index 0000000000000..b31e9b179d26c
--- /dev/null
+++ b/vllm/v1/worker/gpu/block_table.py
@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+
+import torch
+
+from vllm.attention.backends.utils import PAD_SLOT_ID
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+from vllm.v1.utils import CpuGpuBuffer
+
+
+class BlockTables:
+ def __init__(
+ self,
+ block_sizes: list[int],
+ max_num_reqs: int,
+ max_num_batched_tokens: int,
+ max_model_len: int,
+ device: torch.device,
+ pin_memory: bool,
+ ):
+ self.block_sizes = block_sizes
+ self.max_num_reqs = max_num_reqs
+ self.max_num_batched_tokens = max_num_batched_tokens
+ self.max_model_len = max_model_len
+ self.device = device
+ self.pin_memory = pin_memory
+
+ self.num_kv_cache_groups = len(self.block_sizes)
+ # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
+ self.block_tables: list[torch.Tensor] = []
+ for i in range(self.num_kv_cache_groups):
+ block_size = self.block_sizes[i]
+ max_num_blocks = cdiv(self.max_model_len, block_size)
+ block_table = torch.zeros(
+ self.max_num_reqs,
+ max_num_blocks,
+ dtype=torch.int32,
+ device=self.device,
+ )
+ self.block_tables.append(block_table)
+ self.block_table_ptrs = self._make_ptr_tensor(self.block_tables)
+
+ # Block tables used for model's forward pass.
+ # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
+ self.input_block_tables: list[torch.Tensor] = [
+ torch.zeros_like(block_table) for block_table in self.block_tables
+ ]
+ self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
+
+ self.block_table_strides = torch.tensor(
+ [b.stride(0) for b in self.block_tables],
+ dtype=torch.int64,
+ device=self.device,
+ )
+ self.block_sizes_tensor = torch.tensor(
+ self.block_sizes, dtype=torch.int32, device=self.device
+ )
+ self.num_blocks = torch.zeros(
+ self.num_kv_cache_groups,
+ self.max_num_reqs,
+ dtype=torch.int32,
+ device=self.device,
+ )
+ self.slot_mappings = torch.zeros(
+ self.num_kv_cache_groups,
+ self.max_num_batched_tokens,
+ dtype=torch.int64,
+ device=self.device,
+ )
+
+ # Misc buffers.
+ self.req_indices = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
+ self.overwrite = self._make_buffer(self.max_num_reqs, dtype=torch.bool)
+ self.cu_num_new_blocks = self._make_buffer(
+ self.num_kv_cache_groups, self.max_num_reqs + 1, dtype=torch.int32
+ )
+
+ def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
+ return CpuGpuBuffer(
+ *args, dtype=dtype, pin_memory=self.pin_memory, device=self.device
+ )
+
+ def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
+ # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
+ ptrs_tensor_cpu = torch.tensor(
+ [t.data_ptr() for t in x],
+ dtype=torch.uint64,
+ device="cpu",
+ pin_memory=self.pin_memory,
+ )
+ return ptrs_tensor_cpu.to(self.device, non_blocking=True)
+
+ def append_block_ids(
+ self,
+ # [num_reqs]
+ req_indices: list[int],
+ # [num_kv_cache_groups, num_reqs + 1]
+ cu_num_new_blocks: tuple[list[int], ...],
+ # [num_kv_cache_groups, num_new_blocks]
+ new_block_ids: tuple[list[int], ...],
+ # [num_reqs]
+ overwrite: list[bool],
+ ) -> None:
+ num_reqs = len(req_indices)
+ self.req_indices.np[:num_reqs] = req_indices
+ self.overwrite.np[:num_reqs] = overwrite
+ for i in range(self.num_kv_cache_groups):
+ self.cu_num_new_blocks.np[i, : num_reqs + 1] = cu_num_new_blocks[i]
+
+ # NOTE(woosuk): Here, we cannot use a fixed-size buffer because there's
+ # no clear upper bound to the number of new blocks in a single step.
+ # NOTE(woosuk): The buffer has to be cached, because otherwise we cannot
+ # guarantee that the buffer is not freed before the copy is completed.
+ self.new_block_ids_cpu = torch.empty(
+ self.num_kv_cache_groups,
+ max(len(x) for x in new_block_ids),
+ dtype=torch.int32,
+ device="cpu",
+ pin_memory=self.pin_memory,
+ )
+ new_block_ids_np = self.new_block_ids_cpu.numpy()
+ for i in range(self.num_kv_cache_groups):
+ new_block_ids_np[i, : len(new_block_ids[i])] = new_block_ids[i]
+ new_block_ids_gpu = self.new_block_ids_cpu.to(self.device, non_blocking=True)
+
+ _append_block_ids_kernel[(self.num_kv_cache_groups, num_reqs)](
+ self.req_indices.copy_to_gpu(num_reqs),
+ self.cu_num_new_blocks.copy_to_gpu(),
+ self.cu_num_new_blocks.gpu.stride(0),
+ new_block_ids_gpu,
+ new_block_ids_gpu.stride(0),
+ self.overwrite.copy_to_gpu(num_reqs),
+ self.block_table_strides,
+ self.block_table_ptrs,
+ self.num_blocks,
+ self.num_blocks.stride(0),
+ BLOCK_SIZE=1024, # type: ignore
+ )
+
+ def gather_block_tables(
+ self,
+ idx_mapping: torch.Tensor,
+ ) -> tuple[torch.Tensor, ...]:
+ num_reqs = idx_mapping.shape[0]
+ _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs)](
+ idx_mapping,
+ self.block_table_ptrs,
+ self.input_block_table_ptrs,
+ self.block_table_strides,
+ self.num_blocks,
+ self.num_blocks.stride(0),
+ BLOCK_SIZE=1024, # type: ignore
+ )
+ return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
+
+ def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
+ return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
+
+ def compute_slot_mappings(
+ self,
+ query_start_loc: torch.Tensor,
+ positions: torch.Tensor,
+ ) -> torch.Tensor:
+ num_reqs = query_start_loc.shape[0] - 1
+ num_tokens = positions.shape[0]
+ num_groups = self.num_kv_cache_groups
+ _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
+ num_tokens,
+ self.max_num_batched_tokens,
+ query_start_loc,
+ positions,
+ self.input_block_table_ptrs,
+ self.block_table_strides,
+ self.block_sizes_tensor,
+ self.slot_mappings,
+ self.slot_mappings.stride(0),
+ PAD_ID=PAD_SLOT_ID,
+ BLOCK_SIZE=1024, # type: ignore
+ )
+ return self.slot_mappings[:, :num_tokens]
+
+ def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
+ self.slot_mappings.fill_(PAD_SLOT_ID)
+ return self.slot_mappings[:, :num_tokens]
+
+
+@triton.jit
+def _append_block_ids_kernel(
+ # Inputs
+ req_indices, # [num_reqs]
+ cu_num_new_blocks_ptr, # [num_kv_cache_groups, num_reqs + 1]
+ cu_num_new_blocks_stride,
+ new_block_ids_ptr, # [num_kv_cache_groups, num_new_blocks]
+ new_block_ids_stride,
+ overwrite, # [num_reqs]
+ block_table_strides, # [num_kv_cache_groups]
+ # Outputs
+ block_table_ptrs, # [num_kv_cache_groups]
+ num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
+ num_blocks_stride,
+ # Constants
+ BLOCK_SIZE: tl.constexpr,
+):
+ group_id = tl.program_id(0)
+ batch_idx = tl.program_id(1)
+ req_idx = tl.load(req_indices + batch_idx)
+ do_overwrite = tl.load(overwrite + batch_idx)
+
+ group_new_blocks_ptr = cu_num_new_blocks_ptr + group_id * cu_num_new_blocks_stride
+ start_idx = tl.load(group_new_blocks_ptr + batch_idx)
+ end_idx = tl.load(group_new_blocks_ptr + batch_idx + 1)
+ num_new_blocks = end_idx - start_idx
+
+ group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
+ dst_start_idx = tl.load(group_num_blocks_ptr + req_idx) if not do_overwrite else 0
+ dst_end_idx = dst_start_idx + num_new_blocks
+ tl.store(group_num_blocks_ptr + req_idx, dst_end_idx)
+
+ # Destination
+ block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
+ block_table_stride = tl.load(block_table_strides + group_id)
+ row_ptr = block_table_ptr + req_idx * block_table_stride
+
+ group_new_block_ids_ptr = new_block_ids_ptr + group_id * new_block_ids_stride
+ for i in range(0, num_new_blocks, BLOCK_SIZE):
+ offset = i + tl.arange(0, BLOCK_SIZE)
+ block_ids = tl.load(
+ group_new_block_ids_ptr + start_idx + offset, mask=offset < num_new_blocks
+ )
+ tl.store(
+ row_ptr + dst_start_idx + offset, block_ids, mask=offset < num_new_blocks
+ )
+
+
+@triton.jit
+def _gather_block_tables_kernel(
+ batch_idx_to_req_idx, # [batch_size]
+ src_block_table_ptrs, # [num_kv_cache_groups]
+ dst_block_table_ptrs, # [num_kv_cache_groups]
+ block_table_strides, # [num_kv_cache_groups]
+ num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
+ num_blocks_stride,
+ BLOCK_SIZE: tl.constexpr,
+):
+ # kv cache group id
+ group_id = tl.program_id(0)
+ batch_idx = tl.program_id(1)
+ req_idx = tl.load(batch_idx_to_req_idx + batch_idx)
+
+ group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
+ num_blocks = tl.load(group_num_blocks_ptr + req_idx)
+
+ stride = tl.load(block_table_strides + group_id)
+ src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32)
+ src_row_ptr = src_block_table_ptr + req_idx * stride
+ dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32)
+ dst_row_ptr = dst_block_table_ptr + batch_idx * stride
+
+ for i in tl.range(0, num_blocks, BLOCK_SIZE):
+ offset = i + tl.arange(0, BLOCK_SIZE)
+ block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks)
+ tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks)
+
+
+@triton.jit
+def _compute_slot_mappings_kernel(
+ num_tokens,
+ max_num_tokens,
+ cu_num_tokens, # [num_reqs + 1]
+ pos, # [num_tokens]
+ block_table_ptrs, # [num_kv_cache_groups]
+ block_table_strides, # [num_kv_cache_groups]
+ page_sizes, # [num_kv_cache_groups]
+ slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
+ slot_mappings_stride,
+ PAD_ID: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+):
+ # kv cache group id
+ group_id = tl.program_id(0)
+ req_idx = tl.program_id(1)
+ slot_mapping_ptr = slot_mappings_ptr + group_id * slot_mappings_stride
+
+ if req_idx == tl.num_programs(1) - 1:
+ # Pad remaining slots to -1. This is needed for CUDA graphs.
+ for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
+ offset = i + tl.arange(0, BLOCK_SIZE)
+ tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens)
+ return
+
+ block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
+ block_table_stride = tl.load(block_table_strides + group_id)
+ page_size = tl.load(page_sizes + group_id)
+
+ start_idx = tl.load(cu_num_tokens + req_idx)
+ end_idx = tl.load(cu_num_tokens + req_idx + 1)
+ for i in range(start_idx, end_idx, BLOCK_SIZE):
+ offset = i + tl.arange(0, BLOCK_SIZE)
+ positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
+ block_indices = positions // page_size
+ block_numbers = tl.load(
+ block_table_ptr + req_idx * block_table_stride + block_indices
+ )
+ slot_ids = block_numbers * page_size + positions % page_size
+ tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
+
+
+@triton.jit
+def _load_ptr(ptr_to_ptr, elem_dtype):
+ ptr = tl.load(ptr_to_ptr)
+ ptr = tl.cast(ptr, tl.pointer_type(elem_dtype))
+ return tl.multiple_of(ptr, 16)
diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
new file mode 100644
index 0000000000000..ba783e2d0c6fb
--- /dev/null
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import patch
+
+import numpy as np
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
+from vllm.forward_context import set_forward_context
+from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.input_batch import InputBuffers
+
+
+class CudaGraphManager:
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ device: torch.device,
+ ):
+ self.vllm_config = vllm_config
+ self.scheduler_config = vllm_config.scheduler_config
+ self.device = device
+
+ self.max_model_len = vllm_config.model_config.max_model_len
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
+ self.dp_size = vllm_config.parallel_config.data_parallel_size
+ self.compilation_config = vllm_config.compilation_config
+ assert self.compilation_config is not None
+
+ if self.compilation_config.cudagraph_mode is None:
+ self.cudagraph_mode = CUDAGraphMode.NONE
+ else:
+ self.cudagraph_mode = self.compilation_config.cudagraph_mode
+ if self.compilation_config.cudagraph_capture_sizes is not None:
+ cudagraph_sizes = sorted(self.compilation_config.cudagraph_capture_sizes)
+ # Limit the cudagraph sizes to the max decode batch size.
+ self.cudagraph_sizes = [
+ x for x in cudagraph_sizes if x <= self.max_num_reqs
+ ]
+ else:
+ self.cudagraph_sizes = []
+ self.padded_sizes = self._init_padded_sizes()
+
+ self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
+ self.pool = torch.cuda.graph_pool_handle()
+ self.hidden_states: torch.Tensor | None = None
+
+ def _init_padded_sizes(self) -> dict[int, int]:
+ if not self.cudagraph_mode.has_full_cudagraphs():
+ # Full cuda graphs are not used.
+ return {}
+ if not self.cudagraph_sizes:
+ return {}
+
+ padded_sizes: dict[int, int] = {}
+ for i in range(1, self.cudagraph_sizes[-1] + 1):
+ for x in self.cudagraph_sizes:
+ if i <= x:
+ padded_sizes[i] = x
+ break
+ return padded_sizes
+
+ def needs_capture(self) -> bool:
+ return len(self.padded_sizes) > 0
+
+ def get_cudagraph_size(
+ self,
+ scheduler_output: SchedulerOutput,
+ num_tokens_after_padding: int,
+ ) -> int | None:
+ if not self.cudagraph_mode.has_full_cudagraphs():
+ return None
+ if self.cudagraph_mode != CUDAGraphMode.FULL:
+ # TODO(woosuk): Support uniform decode with multiple tokens (spec decoding).
+ all_decode = all(
+ x == 1 for x in scheduler_output.num_scheduled_tokens.values()
+ )
+ if not all_decode:
+ # Prefill is included.
+ return None
+ return self.padded_sizes.get(num_tokens_after_padding)
+
+ def capture_graph(
+ self,
+ batch_size: int,
+ model: nn.Module,
+ input_buffers: InputBuffers,
+ block_tables: BlockTables,
+ attn_metadata_builders: list[AttentionMetadataBuilder],
+ kv_cache_config: KVCacheConfig,
+ ) -> None:
+ assert batch_size not in self.graphs
+
+ # Prepare dummy inputs.
+ input_ids = input_buffers.input_ids.gpu[:batch_size]
+ positions = input_buffers.positions[:batch_size]
+
+ input_buffers.query_start_loc.np[: batch_size + 1] = np.arange(batch_size + 1)
+ input_buffers.query_start_loc.np[batch_size:] = batch_size
+ input_buffers.query_start_loc.copy_to_gpu()
+ # HACK(woosuk): To optimize warmup time, we use 1 (instead of max_model_len)
+ # for seq_lens. This leads to a mismatch between seq_lens (GPU) and
+ # seq_lens_np (CPU), which might cause issues in some attention backends.
+ input_buffers.seq_lens[:batch_size] = 1
+ input_buffers.seq_lens[batch_size:] = 0
+
+ input_block_tables = [x[:batch_size] for x in block_tables.input_block_tables]
+ slot_mappings = block_tables.slot_mappings[:, :batch_size]
+
+ attn_metadata = build_attn_metadata(
+ attn_metadata_builders=attn_metadata_builders,
+ num_reqs=batch_size,
+ num_tokens=batch_size,
+ query_start_loc=input_buffers.query_start_loc,
+ seq_lens=input_buffers.seq_lens,
+ seq_lens_np=np.full(batch_size, self.max_model_len, dtype=np.int32),
+ num_computed_tokens_cpu=None, # FIXME
+ block_tables=input_block_tables,
+ slot_mappings=slot_mappings,
+ kv_cache_config=kv_cache_config,
+ )
+ if self.dp_size > 1:
+ num_tokens_across_dp = torch.full(
+ (self.dp_size,),
+ batch_size,
+ dtype=torch.int32,
+ device="cpu",
+ )
+ else:
+ num_tokens_across_dp = None
+
+ # Warm up.
+ with set_forward_context(
+ attn_metadata,
+ self.vllm_config,
+ num_tokens=batch_size,
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
+ num_tokens_across_dp=num_tokens_across_dp,
+ ):
+ hidden_states = model(
+ input_ids=input_ids,
+ positions=positions,
+ )
+ if self.hidden_states is None:
+ self.hidden_states = torch.empty_like(hidden_states)
+
+ # Capture the graph.
+ graph = torch.cuda.CUDAGraph()
+ with (
+ patch("torch.cuda.empty_cache", lambda: None),
+ set_forward_context(
+ attn_metadata,
+ self.vllm_config,
+ num_tokens=batch_size,
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
+ num_tokens_across_dp=num_tokens_across_dp,
+ ),
+ torch.cuda.graph(graph, self.pool),
+ ):
+ hidden_states = model(
+ input_ids=input_ids,
+ positions=positions,
+ )
+ self.hidden_states[:batch_size] = hidden_states
+ self.graphs[batch_size] = graph
+
+ @torch.inference_mode()
+ def capture(
+ self,
+ model: nn.Module,
+ input_buffers: InputBuffers,
+ block_tables: BlockTables,
+ attn_metadata_builders: list[AttentionMetadataBuilder],
+ kv_cache_config: KVCacheConfig,
+ ) -> None:
+ assert self.needs_capture()
+ # Capture larger graphs first.
+ sizes_to_capture = sorted(self.cudagraph_sizes, reverse=True)
+ if is_global_first_rank():
+ sizes_to_capture = tqdm(sizes_to_capture, desc="Capturing CUDA graphs")
+
+ with graph_capture(device=self.device):
+ for batch_size in sizes_to_capture:
+ self.capture_graph(
+ batch_size,
+ model,
+ input_buffers,
+ block_tables,
+ attn_metadata_builders,
+ kv_cache_config,
+ )
+
+ def run(self, batch_size: int) -> torch.Tensor:
+ assert batch_size in self.graphs
+ self.graphs[batch_size].replay()
+ assert self.hidden_states is not None
+ return self.hidden_states[:batch_size]
diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py
new file mode 100644
index 0000000000000..9bfc7f25bef3a
--- /dev/null
+++ b/vllm/v1/worker/gpu/dp_utils.py
@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.distributed as dist
+
+from vllm.distributed.parallel_state import get_dp_group
+
+
+def get_batch_metadata_across_dp(
+ num_tokens: int,
+ cudagraph_size: int,
+ dp_size: int,
+ dp_rank: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ assert dp_size > 1
+ # Use CPU group to avoid CPU-GPU synchronization.
+ group = get_dp_group().cpu_group
+ tensor = torch.zeros(2, dp_size, dtype=torch.int32, device="cpu")
+ tensor[0][dp_rank] = num_tokens
+ tensor[1][dp_rank] = cudagraph_size
+ dist.all_reduce(tensor, group=group)
+ return tensor[0], tensor[1]
diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py
new file mode 100644
index 0000000000000..7675cb45170b5
--- /dev/null
+++ b/vllm/v1/worker/gpu/input_batch.py
@@ -0,0 +1,440 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from typing import Any
+
+import numba
+import numba.types as types
+import numpy as np
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.utils import random_uuid
+from vllm.utils.math_utils import cdiv
+from vllm.v1.utils import CpuGpuBuffer
+
+
+class InputBuffers:
+ def __init__(
+ self,
+ max_num_reqs: int,
+ max_num_tokens: int,
+ hidden_size: int,
+ vocab_size: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ pin_memory: bool,
+ ):
+ self.max_num_reqs = max_num_reqs
+ self.max_num_tokens = max_num_tokens
+ self.device = device
+ self.pin_memory = pin_memory
+
+ self.idx_mapping = self._make_buffer(max_num_reqs, dtype=torch.int32)
+ self.input_ids = self._make_buffer(max_num_tokens, dtype=torch.int32)
+ self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
+ self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
+ self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+ self.cu_num_logits = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
+
+ # Structured outputs.
+ self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
+ self.grammar_bitmask = self._make_buffer(
+ max_num_reqs, cdiv(vocab_size, 32), dtype=torch.int32
+ )
+
+ def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
+ return CpuGpuBuffer(
+ *args, dtype=dtype, pin_memory=self.pin_memory, device=self.device
+ )
+
+
+@dataclass
+class InputBatch:
+ # batch_idx -> req_id
+ req_ids: list[str]
+ num_reqs: int
+
+ # batch_idx -> req_state_idx
+ idx_mapping: torch.Tensor
+ idx_mapping_np: np.ndarray
+
+ # [num_reqs]
+ # batch_idx -> num_scheduled_tokens
+ num_scheduled_tokens: np.ndarray
+ # sum(num_scheduled_tokens)
+ num_tokens: int
+ num_tokens_after_padding: int
+ num_draft_tokens: int
+
+ # [num_reqs + 1]
+ query_start_loc: torch.Tensor
+ query_start_loc_np: np.ndarray
+ # [num_reqs]
+ seq_lens: torch.Tensor
+ seq_lens_np: np.ndarray
+
+ # [num_tokens_after_padding]
+ input_ids: torch.Tensor
+ # [num_tokens_after_padding]
+ positions: torch.Tensor
+
+ # layer_name -> Metadata
+ attn_metadata: dict[str, Any]
+
+ # [total_num_logits]
+ logits_indices: torch.Tensor
+ # [num_reqs + 1]
+ cu_num_logits: torch.Tensor
+
+ @classmethod
+ def make_dummy(
+ cls,
+ num_reqs: int,
+ num_tokens: int,
+ input_buffers: InputBuffers,
+ device: torch.device,
+ ) -> "InputBatch":
+ assert 0 < num_reqs <= num_tokens
+ req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
+ idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
+ num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
+ num_scheduled_tokens[-1] += num_tokens % num_reqs
+ assert int(num_scheduled_tokens.sum()) == num_tokens
+
+ input_buffers.query_start_loc.np[0] = 0
+ input_buffers.query_start_loc.np[1 : num_reqs + 1] = np.cumsum(
+ num_scheduled_tokens
+ )
+ input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
+ query_start_loc_np = input_buffers.query_start_loc.np[: num_reqs + 1]
+ query_start_loc = input_buffers.query_start_loc.copy_to_gpu()[: num_reqs + 1]
+ # seq_len equals to query_len
+ seq_lens_np = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
+ seq_lens_np[-1] += num_tokens % num_reqs
+ input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
+ input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
+ input_buffers.seq_lens[num_reqs:] = 0
+ seq_lens = input_buffers.seq_lens[:num_reqs]
+
+ input_ids = input_buffers.input_ids.copy_to_gpu(num_tokens)
+ positions = input_buffers.positions[:num_tokens]
+ # attn_metadata = defaultdict(lambda: None)
+ logits_indices = query_start_loc[1:] - 1
+ cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
+ return cls(
+ req_ids=req_ids,
+ num_reqs=num_reqs,
+ idx_mapping=idx_mapping,
+ idx_mapping_np=idx_mapping_np,
+ num_scheduled_tokens=num_scheduled_tokens,
+ num_tokens=num_tokens,
+ num_tokens_after_padding=num_tokens,
+ num_draft_tokens=0,
+ query_start_loc=query_start_loc,
+ query_start_loc_np=query_start_loc_np,
+ seq_lens=seq_lens,
+ seq_lens_np=seq_lens_np,
+ input_ids=input_ids,
+ positions=positions,
+ attn_metadata=None, # type: ignore
+ logits_indices=logits_indices,
+ cu_num_logits=cu_num_logits,
+ )
+
+
+# NOTE: With the type annotations, this function is pre-compiled
+# before the first call.
+@numba.jit(
+ [
+ types.none(
+ types.int32[:], # idx_mapping
+ types.int32[:], # num_scheduled_tokens
+ types.int32[:, :], # prefill_token_ids
+ types.int32[:], # num_computed_prefill_tokens
+ types.int32[:], # prefill_len
+ types.int32[:], # input_ids
+ types.int32[:], # query_start_loc
+ )
+ ],
+ nopython=True,
+ cache=True,
+)
+def _prepare_prefill_inputs(
+ idx_mapping: np.ndarray, # batch_idx -> req_idx
+ num_scheduled_tokens: np.ndarray, # [B]
+ prefill_token_ids: np.ndarray, # [N, max_model_len]
+ num_computed_prefill_tokens: np.ndarray, # [N]
+ prefill_len: np.ndarray, # [N]
+ input_ids: np.ndarray, # [num_input_tokens]
+ query_start_loc: np.ndarray, # [B + 1]
+) -> None:
+ num_reqs = num_scheduled_tokens.shape[0]
+ query_start_loc[0] = 0
+
+ cu_num_tokens = 0
+ for i in range(num_reqs):
+ req_idx = idx_mapping[i]
+ query_len = num_scheduled_tokens[i]
+
+ start = num_computed_prefill_tokens[req_idx]
+ end = min(start + query_len, prefill_len[req_idx])
+ n = end - start
+
+ start_idx = cu_num_tokens
+ input_ids[start_idx : start_idx + n] = prefill_token_ids[req_idx, start:end]
+
+ cu_num_tokens = start_idx + query_len
+ query_start_loc[i + 1] = cu_num_tokens
+
+ # Pad the inputs for CUDA graphs.
+ # Note: pad query_start_loc to be non-decreasing, as kernels
+ # like FlashAttention requires that
+ query_start_loc[num_reqs + 1 :].fill(cu_num_tokens)
+
+
+def prepare_prefill_inputs(
+ idx_mapping: np.ndarray,
+ num_scheduled_tokens: np.ndarray,
+ total_num_tokens: int,
+ prefill_token_ids: np.ndarray,
+ num_computed_prefill_tokens: np.ndarray,
+ prefill_len: np.ndarray,
+ input_ids: CpuGpuBuffer,
+ query_start_loc: CpuGpuBuffer,
+) -> None:
+ _prepare_prefill_inputs(
+ idx_mapping,
+ num_scheduled_tokens,
+ prefill_token_ids,
+ num_computed_prefill_tokens,
+ prefill_len,
+ input_ids.np,
+ query_start_loc.np,
+ )
+ input_ids.copy_to_gpu(total_num_tokens)
+ # NOTE(woosuk): We should copy the whole query_start_loc and seq_lens
+ # tensors from CPU to GPU, because they may include paddings needed
+ # for full CUDA graph mode.
+ query_start_loc.copy_to_gpu()
+
+
+@triton.jit
+def _prepare_pos_seq_lens_kernel(
+ pos_ptr,
+ seq_lens_ptr,
+ idx_mapping_ptr,
+ query_start_loc_ptr,
+ num_computed_tokens_ptr,
+ max_num_reqs,
+ BLOCK_SIZE: tl.constexpr,
+):
+ req_id = tl.program_id(0)
+ num_reqs = tl.num_programs(0) - 1
+ if req_id == num_reqs:
+ # Pad unused seq_lens as 0 for full CUDA graphs.
+ for i in tl.range(num_reqs, max_num_reqs, BLOCK_SIZE):
+ block = i + tl.arange(0, BLOCK_SIZE)
+ mask = block < max_num_reqs
+ tl.store(seq_lens_ptr + block, 0, mask=mask)
+ return
+
+ req_state_idx = tl.load(idx_mapping_ptr + req_id)
+ num_computed_tokens = tl.load(num_computed_tokens_ptr + req_state_idx)
+
+ start = tl.load(query_start_loc_ptr + req_id)
+ end = tl.load(query_start_loc_ptr + req_id + 1)
+ query_len = end - start
+
+ seq_len = num_computed_tokens + query_len
+ tl.store(seq_lens_ptr + req_id, seq_len)
+
+ for i in tl.range(0, query_len, BLOCK_SIZE):
+ block = i + tl.arange(0, BLOCK_SIZE)
+ mask = block < query_len
+ pos = num_computed_tokens + block
+ tl.store(pos_ptr + start + block, pos, mask=mask)
+
+
+def prepare_pos_seq_lens(
+ idx_mapping: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ num_computed_tokens: torch.Tensor,
+ pos: torch.Tensor,
+ seq_lens: torch.Tensor,
+) -> None:
+ num_reqs = idx_mapping.shape[0]
+ # NOTE(woosuk): We do +1 because the last thread block is used
+ # to pad unused seq_lens as 0 for full CUDA graphs.
+ _prepare_pos_seq_lens_kernel[(num_reqs + 1,)](
+ pos,
+ seq_lens,
+ idx_mapping,
+ query_start_loc,
+ num_computed_tokens,
+ seq_lens.shape[0],
+ BLOCK_SIZE=1024,
+ )
+
+
+@triton.jit
+def _combine_sampled_and_draft_tokens_kernel(
+ input_ids_ptr,
+ idx_mapping_ptr,
+ last_sampled_tokens_ptr,
+ query_start_loc_ptr,
+ seq_lens_ptr,
+ prefill_len_ptr,
+ draft_tokens_ptr,
+ draft_tokens_stride,
+ cu_num_logits_ptr,
+ logits_indices_ptr,
+ BLOCK_SIZE: tl.constexpr,
+):
+ batch_idx = tl.program_id(0)
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
+
+ # Get the number of logits and draft tokens.
+ cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx)
+ cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
+ num_logits = cu_num_logits_end - cu_num_logits_start
+ num_draft_tokens = num_logits - 1
+
+ # Compute the logits indices.
+ block = tl.arange(0, BLOCK_SIZE)
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+ logits_start = query_end - num_logits
+ tl.store(
+ logits_indices_ptr + cu_num_logits_start + block,
+ logits_start + block,
+ mask=block < num_logits,
+ )
+
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
+ prefill_len = tl.load(prefill_len_ptr + req_state_idx)
+ if seq_len <= prefill_len:
+ # Handling prefill tokens. No sampled or draft tokens.
+ return
+
+ # Write the last sampled token ID to input_ids.
+ last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
+ tl.store(input_ids_ptr + query_end - num_logits, last_token_id)
+
+ # Write the draft tokens (if any) to input_ids.
+ if num_draft_tokens > 0:
+ mask = block < num_draft_tokens
+ draft_tokens = tl.load(
+ draft_tokens_ptr + req_state_idx * draft_tokens_stride + block,
+ mask=mask,
+ )
+ tl.store(
+ input_ids_ptr + query_end - num_draft_tokens + block,
+ draft_tokens,
+ mask=mask,
+ )
+
+
+def combine_sampled_and_draft_tokens(
+ input_ids: torch.Tensor,
+ idx_mapping: torch.Tensor,
+ last_sampled_tokens: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ seq_lens: torch.Tensor,
+ prefill_len: torch.Tensor,
+ draft_tokens: torch.Tensor,
+ cu_num_logits: torch.Tensor,
+ num_logits: int,
+) -> torch.Tensor:
+ num_reqs = seq_lens.shape[0]
+ num_speculative_steps = draft_tokens.shape[-1]
+
+ logits_indices = torch.empty(
+ num_logits,
+ dtype=torch.int64,
+ device=input_ids.device,
+ )
+ _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
+ input_ids,
+ idx_mapping,
+ last_sampled_tokens,
+ query_start_loc,
+ seq_lens,
+ prefill_len,
+ draft_tokens,
+ draft_tokens.stride(0),
+ cu_num_logits,
+ logits_indices,
+ # NOTE(woosuk): Add 1 to ensure the block can cover the last sampled token
+ # in addition to all draft tokens.
+ BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
+ )
+ return logits_indices
+
+
+@triton.jit
+def _post_update_kernel(
+ idx_mapping_ptr,
+ num_computed_tokens_ptr,
+ last_sampled_tokens_ptr,
+ sampled_tokens_ptr,
+ sampled_tokens_stride,
+ num_sampled_ptr,
+ query_start_loc_ptr,
+ cu_num_logits_ptr,
+):
+ req_id = tl.program_id(0)
+ req_state_idx = tl.load(idx_mapping_ptr + req_id)
+
+ num_sampled = tl.load(num_sampled_ptr + req_id)
+ if num_sampled > 0:
+ token_id = tl.load(
+ sampled_tokens_ptr + req_id * sampled_tokens_stride + num_sampled - 1
+ )
+ tl.store(last_sampled_tokens_ptr + req_state_idx, token_id)
+
+ query_start = tl.load(query_start_loc_ptr + req_id)
+ query_end = tl.load(query_start_loc_ptr + req_id + 1)
+ query_len = query_end - query_start
+
+ num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
+ num_computed += query_len
+ # Consider the rejected tokens in spec decoding.
+ if num_sampled > 0:
+ # NOTE(woosuk): We must skip num_sampled == 0 to account for chunked prefills.
+ logits_start = tl.load(cu_num_logits_ptr + req_id)
+ logits_end = tl.load(cu_num_logits_ptr + req_id + 1)
+ num_logits = logits_end - logits_start
+ num_rejected = num_logits - num_sampled
+ num_computed -= num_rejected
+ tl.store(num_computed_tokens_ptr + req_state_idx, num_computed)
+
+
+def post_update(
+ # [num_reqs]
+ idx_mapping: torch.Tensor,
+ # [max_num_reqs]
+ num_computed_tokens: torch.Tensor,
+ # [max_num_reqs]
+ last_sampled_tokens: torch.Tensor,
+ # [num_reqs, num_speculative_steps + 1]
+ sampled_tokens: torch.Tensor,
+ # [num_reqs]
+ num_sampled: torch.Tensor,
+ # [num_reqs + 1]
+ query_start_loc: torch.Tensor,
+ # [num_reqs + 1]
+ cu_num_logits: torch.Tensor,
+) -> None:
+ num_reqs = idx_mapping.shape[0]
+ _post_update_kernel[(num_reqs,)](
+ idx_mapping,
+ num_computed_tokens,
+ last_sampled_tokens,
+ sampled_tokens,
+ sampled_tokens.stride(0),
+ num_sampled,
+ query_start_loc,
+ cu_num_logits,
+ num_warps=1,
+ )
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
new file mode 100644
index 0000000000000..6e332ee4b75b8
--- /dev/null
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -0,0 +1,919 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import gc
+import time
+from copy import deepcopy
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.utils.mem_constants import GiB_bytes
+from vllm.utils.mem_utils import DeviceMemoryProfiler
+from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import (
+ EMPTY_MODEL_RUNNER_OUTPUT,
+ LogprobsTensors,
+ ModelRunnerOutput,
+)
+from vllm.v1.sample.sampler import SamplerOutput
+from vllm.v1.worker.gpu.async_utils import AsyncOutput, async_barrier
+from vllm.v1.worker.gpu.attn_utils import (
+ build_attn_metadata,
+ get_kv_cache_spec,
+ init_attn_backend,
+ init_kv_cache,
+)
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
+from vllm.v1.worker.gpu.dp_utils import get_batch_metadata_across_dp
+from vllm.v1.worker.gpu.input_batch import (
+ InputBatch,
+ InputBuffers,
+ combine_sampled_and_draft_tokens,
+ post_update,
+ prepare_pos_seq_lens,
+ prepare_prefill_inputs,
+)
+from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
+from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
+from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
+from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
+from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
+
+logger = init_logger(__name__)
+
+
+class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ device: torch.device,
+ ):
+ self.vllm_config = vllm_config
+ self.model_config = vllm_config.model_config
+ self.cache_config = vllm_config.cache_config
+ self.compilation_config = vllm_config.compilation_config
+ self.lora_config = vllm_config.lora_config
+ self.load_config = vllm_config.load_config
+ self.parallel_config = vllm_config.parallel_config
+ self.scheduler_config = vllm_config.scheduler_config
+ self.speculative_config = vllm_config.speculative_config
+ self.observability_config = vllm_config.observability_config
+
+ self.device = device
+ self.pin_memory = is_pin_memory_available()
+ self.dtype = self.model_config.dtype
+ self.kv_cache_dtype = self.dtype
+ if self.cache_config.cache_dtype != "auto":
+ # Quantized KV cache.
+ self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
+ self.cache_config.cache_dtype
+ ]
+ self.is_pooling_model = False
+
+ self.vocab_size = self.model_config.get_vocab_size()
+ self.max_model_len = self.model_config.max_model_len
+ self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
+ self.hidden_size = self.model_config.get_hidden_size()
+
+ self.dp_size = self.parallel_config.data_parallel_size
+ self.dp_rank = self.parallel_config.data_parallel_rank
+
+ self.use_async_scheduling = self.scheduler_config.async_scheduling
+ self.output_copy_stream = torch.cuda.Stream(self.device)
+ self.output_copy_event = torch.cuda.Event()
+ if self.use_async_scheduling:
+ self.input_prep_event = torch.cuda.Event()
+ self.structured_outputs_event = torch.cuda.Event()
+ else:
+ self.input_prep_event = None
+ self.structured_outputs_event = None
+
+ if self.speculative_config is not None:
+ self.do_spec_decode = True
+ self.num_speculative_steps = self.speculative_config.num_speculative_tokens
+ else:
+ self.do_spec_decode = False
+ self.num_speculative_steps = 0
+
+ self.req_states = RequestState(
+ max_num_reqs=self.max_num_reqs,
+ max_model_len=self.max_model_len,
+ max_num_batched_tokens=self.max_num_tokens,
+ num_speculative_steps=self.num_speculative_steps,
+ vocab_size=self.vocab_size,
+ device=self.device,
+ pin_memory=self.pin_memory,
+ )
+ self.input_buffers = InputBuffers(
+ max_num_reqs=self.max_num_reqs,
+ max_num_tokens=self.max_num_tokens,
+ hidden_size=self.hidden_size,
+ vocab_size=self.vocab_size,
+ dtype=self.dtype,
+ device=self.device,
+ pin_memory=self.pin_memory,
+ )
+ self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
+
+ # CUDA graphs.
+ self.cudagraph_manager = CudaGraphManager(
+ vllm_config=self.vllm_config,
+ device=self.device,
+ )
+
+ def get_supported_tasks(self) -> tuple[str]:
+ return ("generate",)
+
+ def load_model(self, *args, **kwargs) -> None:
+ time_before_load = time.perf_counter()
+ with DeviceMemoryProfiler() as m:
+ model_loader = get_model_loader(self.vllm_config.load_config)
+ logger.info("Loading model from scratch...")
+
+ self.model = model_loader.load_model(
+ vllm_config=self.vllm_config,
+ model_config=self.vllm_config.model_config,
+ )
+ if self.lora_config:
+ self.model = self.load_lora_model(
+ self.model,
+ self.vllm_config,
+ self.device,
+ )
+ time_after_load = time.perf_counter()
+
+ self.model_memory_usage = m.consumed_memory
+ logger.info(
+ "Model loading took %.4f GiB and %.6f seconds",
+ m.consumed_memory / GiB_bytes,
+ time_after_load - time_before_load,
+ )
+
+ def get_model(self) -> nn.Module:
+ return self.model
+
+ def get_kv_cache_spec(self):
+ return get_kv_cache_spec(self.vllm_config)
+
+ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+ kv_cache_config = deepcopy(kv_cache_config)
+ self.kv_cache_config = kv_cache_config
+ block_sizes = [
+ kv_cache_group.kv_cache_spec.block_size
+ for kv_cache_group in kv_cache_config.kv_cache_groups
+ ]
+
+ self.block_tables = BlockTables(
+ block_sizes=block_sizes,
+ max_num_reqs=self.max_num_reqs,
+ max_num_batched_tokens=self.max_num_tokens,
+ max_model_len=self.max_model_len,
+ device=self.device,
+ pin_memory=self.pin_memory,
+ )
+
+ self.attn_backends, self.attn_metadata_builders = init_attn_backend(
+ self.kv_cache_config,
+ self.vllm_config,
+ self.device,
+ )
+ # TODO(woosuk): Support other backends.
+ if not all(b.get_name() == "FLASH_ATTN" for b in self.attn_backends.values()):
+ raise NotImplementedError("Only FLASH_ATTN backend is supported currently.")
+
+ self.kv_caches: list[torch.Tensor] = []
+ init_kv_cache(
+ self.kv_caches,
+ self.compilation_config.static_forward_context,
+ self.kv_cache_config,
+ self.attn_backends,
+ self.device,
+ )
+ # Attention groups are not supported.
+ self.attn_groups = [] # type: ignore
+
+ def prepare_dummy_attn_metadata(self, input_batch: InputBatch) -> None:
+ block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
+ slot_mappings = self.block_tables.get_dummy_slot_mappings(
+ input_batch.num_tokens
+ )
+ num_computed_tokens = torch.zeros(
+ input_batch.num_reqs, dtype=torch.int32, device=self.device
+ )
+ attn_metadata = build_attn_metadata(
+ attn_metadata_builders=self.attn_metadata_builders,
+ num_reqs=input_batch.num_reqs,
+ num_tokens=input_batch.num_tokens,
+ query_start_loc=self.input_buffers.query_start_loc,
+ seq_lens=self.input_buffers.seq_lens,
+ seq_lens_np=input_batch.seq_lens_np,
+ num_computed_tokens_cpu=num_computed_tokens,
+ block_tables=block_tables,
+ slot_mappings=slot_mappings,
+ kv_cache_config=self.kv_cache_config,
+ )
+ input_batch.attn_metadata = attn_metadata
+
+ @torch.inference_mode()
+ def _dummy_run(
+ self,
+ num_tokens: int,
+ *args,
+ skip_attn: bool = True,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ num_reqs = min(num_tokens, self.max_num_reqs)
+ input_batch = InputBatch.make_dummy(
+ num_reqs=num_reqs,
+ num_tokens=num_tokens,
+ input_buffers=self.input_buffers,
+ device=self.device,
+ )
+ if not skip_attn:
+ self.prepare_dummy_attn_metadata(input_batch)
+
+ if self.dp_size == 1:
+ num_tokens_across_dp: torch.Tensor | None = None
+ else:
+ num_tokens_across_dp = torch.full(
+ (self.dp_size,), num_tokens, dtype=torch.int32, device="cpu"
+ )
+ num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)
+ with (
+ self.maybe_dummy_run_with_lora(
+ self.lora_config,
+ input_batch.num_scheduled_tokens,
+ num_sampled_tokens,
+ ),
+ set_forward_context(
+ input_batch.attn_metadata,
+ self.vllm_config,
+ num_tokens=num_tokens,
+ num_tokens_across_dp=num_tokens_across_dp,
+ ),
+ ):
+ hidden_states = self.model(
+ input_ids=input_batch.input_ids,
+ positions=input_batch.positions,
+ )
+ sample_hidden_states = hidden_states[input_batch.logits_indices]
+ return hidden_states, sample_hidden_states
+
+ @torch.inference_mode()
+ def _dummy_sampler_run(
+ self,
+ hidden_states: torch.Tensor,
+ ) -> None:
+ num_reqs = hidden_states.shape[0]
+ sampling_metadata = SamplingMetadata.make_dummy(
+ num_reqs=num_reqs,
+ device=self.device,
+ )
+ logits = self.model.compute_logits(hidden_states)
+ self.sampler(logits, sampling_metadata)
+
+ @torch.inference_mode()
+ def profile_run(self) -> None:
+ hidden_states, sample_hidden_states = self._dummy_run(
+ self.max_num_tokens,
+ skip_attn=True,
+ )
+ self._dummy_sampler_run(sample_hidden_states)
+ torch.cuda.synchronize()
+ del hidden_states, sample_hidden_states
+ gc.collect()
+
+ def reset_mm_cache(self) -> None:
+ pass
+
+ def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:
+ # SP is not supported yet.
+ return num_scheduled_tokens
+
+ @torch.inference_mode()
+ def capture_model(self) -> int:
+ if not self.cudagraph_manager.needs_capture():
+ logger.warning(
+ "Skipping CUDA graph capture. To turn on CUDA graph capture, "
+ "ensure `cudagraph_mode` was not manually set to `NONE`"
+ )
+ return 0
+
+ start_time = time.perf_counter()
+ gc.collect()
+ torch.cuda.empty_cache()
+ start_free_gpu_memory = torch.cuda.mem_get_info()[0]
+
+ with self.maybe_setup_dummy_loras(self.lora_config):
+ self.cudagraph_manager.capture(
+ model=self.model,
+ input_buffers=self.input_buffers,
+ block_tables=self.block_tables,
+ attn_metadata_builders=self.attn_metadata_builders,
+ kv_cache_config=self.kv_cache_config,
+ )
+
+ end_time = time.perf_counter()
+ end_free_gpu_memory = torch.cuda.mem_get_info()[0]
+ elapsed_time = end_time - start_time
+ cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
+ # This usually takes 5~20 seconds.
+ logger.info(
+ "Graph capturing finished in %.0f secs, took %.2f GiB",
+ elapsed_time,
+ cuda_graph_size / (1 << 30),
+ )
+ return cuda_graph_size
+
+ def warmup_for_prefill(self) -> None:
+ # For FlashInfer, we would like to execute a dummy prefill run
+ # to trigger JIT compilation.
+ if all("FLASHINFER" in b.get_name() for b in self.attn_backends.values()):
+ self._dummy_run(self.max_num_tokens, skip_attn=False)
+ torch.cuda.synchronize()
+
+ def update_states(self, scheduler_output: SchedulerOutput) -> None:
+ if scheduler_output.preempted_req_ids is not None:
+ for req_id in scheduler_output.preempted_req_ids:
+ self.req_states.remove_request(req_id)
+ for req_id in scheduler_output.finished_req_ids:
+ self.req_states.remove_request(req_id)
+
+ # TODO(woosuk): Change SchedulerOutput.
+ req_indices: list[int] = []
+ cu_num_new_blocks = tuple(
+ [0] for _ in range(self.block_tables.num_kv_cache_groups)
+ )
+ new_block_ids: tuple[list[int], ...] = tuple(
+ [] for _ in range(self.block_tables.num_kv_cache_groups)
+ )
+ overwrite: list[bool] = []
+
+ # Add new requests.
+ for new_req_data in scheduler_output.scheduled_new_reqs:
+ assert new_req_data.prompt_token_ids is not None
+ assert new_req_data.prefill_token_ids is not None
+ assert new_req_data.sampling_params is not None
+ req_id = new_req_data.req_id
+ self.req_states.add_request(
+ req_id=req_id,
+ prompt_len=len(new_req_data.prompt_token_ids),
+ prefill_token_ids=new_req_data.prefill_token_ids,
+ num_computed_tokens=new_req_data.num_computed_tokens,
+ sampling_params=new_req_data.sampling_params,
+ lora_request=new_req_data.lora_request,
+ )
+
+ req_index = self.req_states.req_id_to_index[req_id]
+ req_indices.append(req_index)
+ for i, block_ids in enumerate(new_req_data.block_ids):
+ x = cu_num_new_blocks[i][-1]
+ cu_num_new_blocks[i].append(x + len(block_ids))
+ new_block_ids[i].extend(block_ids)
+ overwrite.append(True)
+ # Update the GPU tensors for request states.
+ if scheduler_output.scheduled_new_reqs:
+ self.req_states.prefill_len.copy_to_gpu()
+
+ # Add new blocks for the existing requests.
+ cached_reqs = scheduler_output.scheduled_cached_reqs
+ for i, req_id in enumerate(cached_reqs.req_ids):
+ req_index = self.req_states.req_id_to_index[req_id]
+
+ req_new_block_ids = cached_reqs.new_block_ids[i]
+ if req_new_block_ids is not None:
+ req_indices.append(req_index)
+ for group_id, block_ids in enumerate(req_new_block_ids):
+ x = cu_num_new_blocks[group_id][-1]
+ cu_num_new_blocks[group_id].append(x + len(block_ids))
+ new_block_ids[group_id].extend(block_ids)
+ overwrite.append(False)
+
+ if req_indices:
+ self.block_tables.append_block_ids(
+ req_indices=req_indices,
+ cu_num_new_blocks=cu_num_new_blocks,
+ new_block_ids=new_block_ids,
+ overwrite=overwrite,
+ )
+
+ def prepare_inputs(
+ self,
+ scheduler_output: SchedulerOutput,
+ num_tokens_after_padding: int,
+ ) -> InputBatch:
+ num_tokens = scheduler_output.total_num_scheduled_tokens
+ assert num_tokens > 0
+ num_reqs = len(scheduler_output.num_scheduled_tokens)
+
+ # Decode first, then prefill.
+ # batch_idx -> req_id
+ req_ids = sorted(
+ scheduler_output.num_scheduled_tokens.keys(),
+ key=lambda k: scheduler_output.num_scheduled_tokens[k],
+ )
+ num_scheduled_tokens = np.array(
+ [scheduler_output.num_scheduled_tokens[i] for i in req_ids], dtype=np.int32
+ )
+
+ idx_mapping_list = [
+ self.req_states.req_id_to_index[req_id] for req_id in req_ids
+ ]
+ idx_mapping = self.input_buffers.idx_mapping
+ idx_mapping.np[:num_reqs] = idx_mapping_list
+ idx_mapping_np = idx_mapping.np[:num_reqs]
+ idx_mapping = idx_mapping.copy_to_gpu(num_reqs)
+
+ # Get the number of draft tokens for each request.
+ if not scheduler_output.scheduled_spec_decode_tokens:
+ # No draft token scheduled (common case).
+ total_num_draft_tokens = 0
+ total_num_logits = num_reqs
+ cu_num_logits = torch.arange(
+ num_reqs + 1, device=self.device, dtype=torch.int32
+ )
+ else:
+ draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+ num_draft_tokens = np.array(
+ [
+ len(draft_tokens[req_id]) if req_id in draft_tokens else 0
+ for req_id in req_ids
+ ],
+ dtype=np.int32,
+ )
+ total_num_draft_tokens = int(num_draft_tokens.sum())
+ total_num_logits = num_reqs + total_num_draft_tokens
+
+ np.cumsum(
+ num_draft_tokens + 1,
+ out=self.input_buffers.cu_num_logits.np[1 : num_reqs + 1],
+ )
+ cu_num_logits = self.input_buffers.cu_num_logits.copy_to_gpu(num_reqs + 1)
+
+ # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
+ block_tables = self.block_tables.gather_block_tables(idx_mapping)
+
+ # Copy prefill tokens from CPU to GPU and get query_start_loc.
+ prepare_prefill_inputs(
+ idx_mapping_np,
+ num_scheduled_tokens,
+ num_tokens,
+ self.req_states.prefill_token_ids,
+ self.req_states.num_computed_prefill_tokens,
+ self.req_states.prefill_len.np,
+ self.input_buffers.input_ids,
+ self.input_buffers.query_start_loc,
+ )
+ query_start_loc = self.input_buffers.query_start_loc
+ query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
+ query_start_loc_np = query_start_loc.np[: num_reqs + 1]
+
+ # Prepare positions and seq_lens.
+ prepare_pos_seq_lens(
+ idx_mapping,
+ query_start_loc_gpu,
+ self.req_states.num_computed_tokens,
+ self.input_buffers.positions,
+ self.input_buffers.seq_lens,
+ )
+ seq_lens = self.input_buffers.seq_lens[:num_reqs]
+
+ # Some input token ids are directly read from the last sampled tokens
+ # and draft tokens. Also, get the logits indices to sample tokens from.
+ logits_indices = combine_sampled_and_draft_tokens(
+ self.input_buffers.input_ids.gpu,
+ idx_mapping,
+ self.req_states.last_sampled_tokens,
+ query_start_loc_gpu,
+ seq_lens,
+ self.req_states.prefill_len.gpu,
+ self.req_states.draft_tokens,
+ cu_num_logits,
+ total_num_logits,
+ )
+
+ # Compute slot mappings: [num_kv_cache_groups, num_tokens]
+ slot_mappings = self.block_tables.compute_slot_mappings(
+ query_start_loc_gpu, self.input_buffers.positions[:num_tokens]
+ )
+
+ # Get num_computed_tokens.
+ # HACK(woosuk): Here, we use num_computed_tokens on GPU instead of
+ # num_computed_tokens_cpu. This works for most cases.
+ num_computed_tokens = self.req_states.num_computed_tokens[idx_mapping]
+ # HACK(woosuk): Only GPU has the exact seq_lens because at this point
+ # CPU does not know how many draft tokens are accepted/rejected in the
+ # previous step. Therefore, we use max_model_len to be safe.
+ # NOTE(woosuk): This only works for FA3 backend.
+ seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
+
+ # Layer name -> attention metadata.
+ attn_metadata = build_attn_metadata(
+ attn_metadata_builders=self.attn_metadata_builders,
+ num_reqs=num_reqs,
+ num_tokens=num_tokens,
+ query_start_loc=self.input_buffers.query_start_loc,
+ seq_lens=self.input_buffers.seq_lens,
+ seq_lens_np=seq_lens_np,
+ num_computed_tokens_cpu=num_computed_tokens,
+ block_tables=block_tables,
+ slot_mappings=slot_mappings,
+ kv_cache_config=self.kv_cache_config,
+ )
+
+ input_ids = self.input_buffers.input_ids.gpu[:num_tokens_after_padding]
+ positions = self.input_buffers.positions[:num_tokens_after_padding]
+ return InputBatch(
+ req_ids=req_ids,
+ num_reqs=num_reqs,
+ idx_mapping=idx_mapping,
+ idx_mapping_np=idx_mapping_np,
+ num_scheduled_tokens=num_scheduled_tokens,
+ num_tokens=num_tokens,
+ num_tokens_after_padding=num_tokens_after_padding,
+ num_draft_tokens=total_num_draft_tokens,
+ query_start_loc=query_start_loc_gpu,
+ query_start_loc_np=query_start_loc_np,
+ seq_lens=seq_lens,
+ seq_lens_np=seq_lens_np,
+ input_ids=input_ids,
+ positions=positions,
+ attn_metadata=attn_metadata,
+ logits_indices=logits_indices,
+ cu_num_logits=cu_num_logits,
+ )
+
+ def sample(
+ self,
+ hidden_states: torch.Tensor,
+ input_batch: InputBatch,
+ sampling_metadata: SamplingMetadata,
+ grammar_output: GrammarOutput | None,
+ ) -> tuple[SamplerOutput, torch.Tensor]:
+ sample_hidden_states = hidden_states[input_batch.logits_indices]
+ logits = self.model.compute_logits(sample_hidden_states)
+ if grammar_output is not None:
+ # Apply grammar bitmask to the logits in-place.
+ # TODO(woosuk): Make compatible with spec decoding.
+ assert input_batch.num_draft_tokens == 0
+ with async_barrier(self.structured_outputs_event):
+ apply_grammar_bitmask(
+ logits,
+ input_batch.req_ids,
+ grammar_output.structured_output_request_ids,
+ grammar_output.grammar_bitmask,
+ self.input_buffers,
+ )
+
+ # Sample tokens and compute logprobs (if needed).
+ sampler_output = self.sampler(logits, sampling_metadata)
+
+ # Get the number of sampled tokens.
+ prefill_len = self.req_states.prefill_len.gpu[input_batch.idx_mapping]
+ is_chunked_prefilling = input_batch.seq_lens < prefill_len
+ if input_batch.num_draft_tokens == 0:
+ # No draft tokens (common case).
+ # 0 if chunked-prefilling, 1 if not.
+ num_sampled = (~is_chunked_prefilling).int()
+ else:
+ # Draft tokens for spec decoding.
+ input_ids = input_batch.input_ids[input_batch.logits_indices]
+ sampled_tokens, num_sampled = rejection_sample(
+ sampler_output.sampled_token_ids,
+ input_ids,
+ input_batch.cu_num_logits,
+ self.num_speculative_steps,
+ )
+ num_sampled *= ~is_chunked_prefilling
+ sampler_output.sampled_token_ids = sampled_tokens
+ # TODO(woosuk): Support logprobs with spec decoding.
+ return sampler_output, num_sampled
+
+ def compute_prompt_logprobs(
+ self,
+ hidden_states: torch.Tensor,
+ input_batch: InputBatch,
+ ) -> dict[str, LogprobsTensors]:
+ idx_mapping_np = input_batch.idx_mapping_np
+ needs_prompt_logprobs = self.req_states.needs_prompt_logprobs[idx_mapping_np]
+ if not np.any(needs_prompt_logprobs):
+ # No request asks for prompt logprobs.
+ return {}
+
+ prompt_lens = self.req_states.prompt_len[idx_mapping_np]
+ # NOTE(woosuk): -1 because the last prompt token's hidden state is not
+ # needed for prompt logprobs.
+ computed_prefill = self.req_states.num_computed_prefill_tokens[idx_mapping_np]
+ includes_prompt = computed_prefill < prompt_lens - 1
+ # NOTE(woosuk): If the request was resumed after preemption, its prompt
+ # logprobs must have been computed before preemption. Skip.
+ resumed_after_prompt = (
+ prompt_lens < self.req_states.prefill_len.np[idx_mapping_np]
+ )
+ needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt
+ if not np.any(needs_prompt_logprobs):
+ return {}
+
+ # Just to be safe, clone the input ids.
+ n = input_batch.num_tokens
+ # Shift the input ids by one.
+ token_ids = torch.empty_like(input_batch.input_ids[:n])
+ token_ids[: n - 1] = input_batch.input_ids[1:n]
+ # To avoid out-of-bound access, set the last token id to 0.
+ token_ids[n - 1] = 0
+
+ # Handle chunked prompts.
+ pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
+ is_prompt_chunked = pos_after_step < prompt_lens
+ prefill_token_ids = self.req_states.prefill_token_ids
+ query_start_loc = self.input_buffers.query_start_loc.np
+ for i, req_id in enumerate(input_batch.req_ids):
+ if not needs_prompt_logprobs[i]:
+ continue
+ if not is_prompt_chunked[i]:
+ continue
+ # The prompt is chunked. Get the next prompt token.
+ req_idx = input_batch.idx_mapping_np[i]
+ next_prompt_token = int(prefill_token_ids[req_idx, pos_after_step[i]])
+ idx = int(query_start_loc[i + 1] - 1)
+ # Set the next prompt token.
+ # NOTE(woosuk): This triggers a GPU operation.
+ token_ids[idx] = next_prompt_token
+
+ # NOTE(woosuk): We mask out logprobs for negative tokens.
+ prompt_logprobs, prompt_ranks = compute_prompt_logprobs(
+ token_ids,
+ hidden_states[:n],
+ self.model.compute_logits,
+ )
+
+ prompt_token_ids = token_ids.unsqueeze(-1)
+ prompt_logprobs_dict: dict[str, LogprobsTensors] = {}
+ for i, req_id in enumerate(input_batch.req_ids):
+ if not needs_prompt_logprobs[i]:
+ continue
+
+ start_idx = query_start_loc[i]
+ end_idx = query_start_loc[i + 1]
+ assert start_idx < end_idx, (
+ f"start_idx ({start_idx}) >= end_idx ({end_idx})"
+ )
+ logprobs = LogprobsTensors(
+ logprob_token_ids=prompt_token_ids[start_idx:end_idx],
+ logprobs=prompt_logprobs[start_idx:end_idx],
+ selected_token_ranks=prompt_ranks[start_idx:end_idx],
+ )
+
+ req_extra_data = self.req_states.extra_data[req_id]
+ prompt_logprobs_list = req_extra_data.in_progress_prompt_logprobs
+ if is_prompt_chunked[i]:
+ # Prompt is chunked. Do not return the logprobs yet.
+ prompt_logprobs_list.append(logprobs)
+ continue
+
+ if prompt_logprobs_list:
+ # Merge the in-progress logprobs.
+ prompt_logprobs_list.append(logprobs)
+ logprobs = LogprobsTensors(
+ logprob_token_ids=torch.cat(
+ [x.logprob_token_ids for x in prompt_logprobs_list]
+ ),
+ logprobs=torch.cat([x.logprobs for x in prompt_logprobs_list]),
+ selected_token_ranks=torch.cat(
+ [x.selected_token_ranks for x in prompt_logprobs_list]
+ ),
+ )
+ prompt_logprobs_list.clear()
+
+ prompt_logprobs_dict[req_id] = logprobs
+ return prompt_logprobs_dict
+
+ def postprocess(
+ self,
+ input_batch: InputBatch,
+ sampled_tokens: torch.Tensor,
+ num_sampled: torch.Tensor,
+ ) -> None:
+ # Update the number of computed tokens.
+ post_update(
+ input_batch.idx_mapping,
+ self.req_states.num_computed_tokens,
+ self.req_states.last_sampled_tokens,
+ sampled_tokens,
+ num_sampled,
+ input_batch.query_start_loc,
+ input_batch.cu_num_logits,
+ )
+
+ # Update the number of computed prefill tokens.
+ idx_mapping_np = input_batch.idx_mapping_np
+ computed_prefill = self.req_states.num_computed_prefill_tokens
+ # TODO(woosuk): Simplify this.
+ computed_prefill[idx_mapping_np] = np.minimum(
+ computed_prefill[idx_mapping_np] + input_batch.num_scheduled_tokens,
+ self.req_states.prefill_len.np[idx_mapping_np],
+ )
+
+ def get_cudagraph_and_dp_padding(
+ self,
+ scheduler_output: SchedulerOutput,
+ ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
+ total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+ if self.dp_size == 1:
+ # No DP. Only consider CUDA graphs.
+ if total_num_scheduled_tokens == 0:
+ # Special case: no tokens to run.
+ return CUDAGraphMode.NONE, 0, None
+
+ cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
+ scheduler_output, total_num_scheduled_tokens
+ )
+ if cudagraph_size is not None:
+ # Use full CUDA graph.
+ return CUDAGraphMode.FULL, cudagraph_size, None
+ # Fall back to eager mode.
+ # TODO(woosuk): Support piecewise CUDA graphs.
+ return CUDAGraphMode.NONE, total_num_scheduled_tokens, None
+
+ # Consider DP padding and CUDA graph.
+ if total_num_scheduled_tokens == 0:
+ # Special handling is needed for 0.
+ cudagraph_size_before_dp: int | None = 0
+ else:
+ cudagraph_size_before_dp = self.cudagraph_manager.get_cudagraph_size(
+ scheduler_output, total_num_scheduled_tokens
+ )
+ if cudagraph_size_before_dp is None:
+ cudagraph_size_before_dp = -1
+
+ assert cudagraph_size_before_dp is not None
+ num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(
+ total_num_scheduled_tokens,
+ cudagraph_size_before_dp,
+ self.dp_size,
+ self.dp_rank,
+ )
+ if all(cudagraph_size_across_dp >= 0):
+ # If all ranks can use CUDA graph, pad to the maximum number of tokens
+ # across DP and use CUDA graph.
+ num_tokens_after_padding = int(cudagraph_size_across_dp.max().item())
+ cudagraph_mode = CUDAGraphMode.FULL
+ else:
+ # If any of the ranks cannot use CUDA graph, use eager mode for all ranks.
+ # No padding is needed except for ranks that have no tokens to run.
+ num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
+ num_tokens_after_padding = num_tokens_across_dp[self.dp_rank]
+ cudagraph_mode = CUDAGraphMode.NONE
+ return cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp
+
+ @torch.inference_mode()
+ def execute_model(
+ self,
+ scheduler_output: SchedulerOutput,
+ intermediate_tensors: Any | None = None,
+ dummy_run: bool = False,
+ ) -> ModelRunnerOutput | None:
+ assert intermediate_tensors is None
+ if scheduler_output.total_num_scheduled_tokens == 0 and not dummy_run:
+ # No need to run the model.
+ with async_barrier(self.input_prep_event):
+ self.update_states(scheduler_output)
+ return EMPTY_MODEL_RUNNER_OUTPUT
+
+ # NOTE: Call this before the async barrier so CPU all-reduce and
+ # GPU execution can overlap.
+ cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp = (
+ self.get_cudagraph_and_dp_padding(scheduler_output)
+ )
+ with async_barrier(self.input_prep_event):
+ self.update_states(scheduler_output)
+ if num_tokens_after_padding == 0:
+ # All DP ranks have zero tokens to run.
+ return EMPTY_MODEL_RUNNER_OUTPUT
+
+ if not dummy_run:
+ # Common case.
+ # Prepare all the inputs and copy to the input buffers.
+ input_batch = self.prepare_inputs(
+ scheduler_output,
+ num_tokens_after_padding,
+ )
+
+ # NOTE(woosuk): Sampling metadata should be built under the async
+ # barrier to avoid race conditions.
+ pos = input_batch.positions[input_batch.logits_indices]
+ sampling_metadata = self.req_states.make_sampling_metadata(
+ input_batch.idx_mapping_np, pos
+ )
+ if input_batch.num_draft_tokens > 0:
+ sampling_metadata = self.req_states.expand_sampling_metadata(
+ sampling_metadata, input_batch.cu_num_logits
+ )
+
+ if self.lora_config:
+ # Activate LoRA adapters.
+ lora_inputs = self.req_states.make_lora_inputs(
+ input_batch.req_ids,
+ input_batch.idx_mapping_np,
+ input_batch.num_scheduled_tokens,
+ )
+ self._set_active_loras(*lora_inputs)
+ else:
+ # No actual tokens to run. A dummy run for DP.
+ num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
+ input_batch = InputBatch.make_dummy(
+ num_reqs=num_reqs,
+ num_tokens=num_tokens_after_padding,
+ input_buffers=self.input_buffers,
+ device=self.device,
+ )
+ self.prepare_dummy_attn_metadata(input_batch)
+ sampling_metadata = None
+
+ # Run model.
+ if cudagraph_mode == CUDAGraphMode.FULL:
+ # Run CUDA graph.
+ # NOTE(woosuk): Here, we don't need to pass the input tensors,
+ # because they are already copied to the CUDA graph input buffers.
+ hidden_states = self.cudagraph_manager.run(
+ input_batch.num_tokens_after_padding
+ )
+ else:
+ # Run PyTorch model in eager mode.
+ # TODO(woosuk): Support piecewise CUDA graph.
+ with set_forward_context(
+ input_batch.attn_metadata,
+ self.vllm_config,
+ num_tokens=input_batch.num_tokens_after_padding,
+ cudagraph_runtime_mode=cudagraph_mode,
+ num_tokens_across_dp=num_tokens_across_dp,
+ ):
+ hidden_states = self.model(
+ input_ids=input_batch.input_ids,
+ positions=input_batch.positions,
+ )
+
+ self.execute_model_state = hidden_states, input_batch, sampling_metadata
+ return None
+
+ @torch.inference_mode()
+ def sample_tokens(
+ self,
+ grammar_output: GrammarOutput | None,
+ ) -> AsyncOutput | ModelRunnerOutput:
+ assert self.execute_model_state is not None
+ hidden_states, input_batch, sampling_metadata = self.execute_model_state
+ self.execute_model_state = None # type: ignore
+ assert sampling_metadata is not None
+
+ sampler_output, num_sampled_tokens = self.sample(
+ hidden_states, input_batch, sampling_metadata, grammar_output
+ )
+ prompt_logprobs_dict = self.compute_prompt_logprobs(hidden_states, input_batch)
+
+ # Prepare the model runner output.
+ model_runner_output = ModelRunnerOutput(
+ req_ids=input_batch.req_ids,
+ # NOTE(woosuk): req_id_to_index is unused in this model runner.
+ # Only for compatibility with the existing model runner and scheduler.
+ req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+ sampled_token_ids=None, # type: ignore
+ logprobs=None,
+ prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore
+ pooler_output=[],
+ kv_connector_output=None,
+ num_nans_in_logits=None,
+ )
+ async_output = AsyncOutput(
+ model_runner_output=model_runner_output,
+ sampler_output=sampler_output,
+ num_sampled_tokens=num_sampled_tokens,
+ copy_stream=self.output_copy_stream,
+ copy_event=self.output_copy_event,
+ )
+
+ # Postprocess results and update request states.
+ # NOTE: This is intentionally done after creating the AsyncOutput,
+ # ensuring that `copy_event` is recorded before calling postprocess.
+ # This sequencing may slightly reduce latency as async D2H copy does not
+ # need to wait for the postprocess to finish.
+ self.postprocess(
+ input_batch, sampler_output.sampled_token_ids, num_sampled_tokens
+ )
+
+ if self.use_async_scheduling:
+ return async_output
+ return async_output.get_output()
diff --git a/vllm/v1/worker/gpu/sampler.py b/vllm/v1/worker/gpu/sampler.py
new file mode 100644
index 0000000000000..c48ed2d8ca167
--- /dev/null
+++ b/vllm/v1/worker/gpu/sampler.py
@@ -0,0 +1,329 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
+
+import torch
+
+from vllm.config.model import LogprobsMode
+from vllm.triton_utils import tl, triton
+from vllm.v1.outputs import LogprobsTensors, SamplerOutput
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.worker.gpu.states import SamplingMetadata
+
+
+class Sampler:
+ def __init__(
+ self,
+ logprobs_mode: LogprobsMode = "raw_logprobs",
+ ):
+ if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
+ raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")
+ self.logprobs_mode = logprobs_mode
+
+ def __call__(
+ self,
+ logits: torch.Tensor,
+ sampling_metadata: SamplingMetadata,
+ ) -> SamplerOutput:
+ if sampling_metadata.max_num_logprobs is not None:
+ if self.logprobs_mode == "processed_logprobs":
+ sampled, logits = self.sample(
+ logits, sampling_metadata, return_logits=True
+ )
+ else:
+ assert self.logprobs_mode == "raw_logprobs"
+ sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
+
+ logprobs_tensors = compute_topk_logprobs(
+ logits,
+ sampling_metadata.max_num_logprobs,
+ sampled,
+ )
+ else:
+ sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
+ logprobs_tensors = None
+
+ # These are GPU tensors.
+ sampler_output = SamplerOutput(
+ # The sampled tokens are expanded to 2D tensor with shape
+ # [num_requests, 1], where each row represents one generated
+ # token per request.
+ sampled_token_ids=sampled.view(-1, 1),
+ logprobs_tensors=logprobs_tensors,
+ )
+ return sampler_output
+
+ def sample(
+ self,
+ logits: torch.Tensor,
+ sampling_metadata: SamplingMetadata,
+ return_logits: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ is_greedy = sampling_metadata.temperature == 0
+ temp = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
+ logits = logits / temp.view(-1, 1)
+ logits = apply_top_k_top_p(
+ logits, sampling_metadata.top_k, sampling_metadata.top_p
+ )
+
+ sampled = gumbel_sample(
+ logits,
+ sampling_metadata.temperature,
+ sampling_metadata.seeds,
+ sampling_metadata.pos,
+ apply_temperature=False,
+ )
+ return sampled, logits if return_logits else None
+
+
+@triton.jit
+def _gumbel_sample_kernel(
+ local_argmax_ptr,
+ local_argmax_stride,
+ local_max_ptr,
+ local_max_stride,
+ logits_ptr,
+ logits_stride,
+ seeds_ptr,
+ pos_ptr,
+ temp_ptr,
+ vocab_size,
+ BLOCK_SIZE: tl.constexpr,
+ APPLY_TEMPERATURE: tl.constexpr,
+):
+ req_idx = tl.program_id(0)
+ block_idx = tl.program_id(1)
+ block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ mask = block < vocab_size
+ logits = tl.load(
+ logits_ptr + req_idx * logits_stride + block,
+ mask=mask,
+ other=float("-inf"),
+ )
+
+ temp = tl.load(temp_ptr + req_idx)
+ if temp != 0.0:
+ # Calculate the seed for gumbel noise.
+ seed = tl.load(seeds_ptr + req_idx)
+ pos = tl.load(pos_ptr + req_idx)
+ gumbel_seed = tl.randint(seed, pos)
+
+ # Generate gumbel noise.
+ r = tl.rand(gumbel_seed, block).to(tl.float64)
+ gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
+ gumbel_noise = gumbel_noise.to(tl.float32)
+
+ # Apply temperature.
+ if APPLY_TEMPERATURE:
+ # NOTE(woosuk): Use div_rn to match the behavior of torch.
+ logits = tl.div_rn(logits, temp.to(tl.float32))
+
+ # Apply gumbel noise.
+ logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
+
+ idx = tl.argmax(logits, axis=0)
+ token_id = block_idx * BLOCK_SIZE + idx
+ value = tl.max(logits, axis=0)
+ tl.store(local_argmax_ptr + req_idx * local_argmax_stride + block_idx, token_id)
+ tl.store(local_max_ptr + req_idx * local_max_stride + block_idx, value)
+
+
+def gumbel_sample(
+ logits: torch.Tensor, # [num_reqs, vocab_size]
+ temperature: torch.Tensor, # [num_reqs]
+ seed: torch.Tensor, # [num_reqs]
+ pos: torch.Tensor, # [num_reqs]
+ apply_temperature: bool,
+) -> torch.Tensor:
+ num_reqs, vocab_size = logits.shape
+ BLOCK_SIZE = 1024
+ num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+ local_argmax = torch.empty(
+ num_reqs,
+ num_blocks,
+ dtype=torch.int64,
+ device=logits.device,
+ )
+ local_max = torch.empty(
+ num_reqs,
+ num_blocks,
+ dtype=torch.float32,
+ device=logits.device,
+ )
+ _gumbel_sample_kernel[(num_reqs, num_blocks)](
+ local_argmax,
+ local_argmax.stride(0),
+ local_max,
+ local_max.stride(0),
+ logits,
+ logits.stride(0),
+ seed,
+ pos,
+ temperature,
+ vocab_size,
+ BLOCK_SIZE=BLOCK_SIZE,
+ APPLY_TEMPERATURE=apply_temperature,
+ )
+ # NOTE(woosuk): Use int64 for later indexing.
+ max_block_idx = local_max.argmax(dim=-1, keepdim=True)
+ sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
+ return sampled
+
+
+@triton.jit
+def _topk_log_softmax_kernel(
+ output_ptr,
+ logits_ptr,
+ logits_stride,
+ topk_ids_ptr,
+ topk,
+ vocab_size,
+ BLOCK_SIZE: tl.constexpr,
+ PADDED_TOPK: tl.constexpr,
+):
+ req_idx = tl.program_id(0)
+ row_ptr = logits_ptr + req_idx * logits_stride
+
+ max_val = float("-inf")
+ for i in range(0, vocab_size, BLOCK_SIZE):
+ block = i + tl.arange(0, BLOCK_SIZE)
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
+ max_val = tl.max(tl.maximum(logits, max_val))
+ max_val = max_val.to(tl.float32) # type: ignore
+
+ se = 0.0
+ for i in range(0, vocab_size, BLOCK_SIZE):
+ block = i + tl.arange(0, BLOCK_SIZE)
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=0.0)
+ # NOTE(woosuk): Make sure that logits and all following operations use FP32.
+ logits = logits.to(tl.float32)
+ e = tl.exp(logits - max_val)
+ e = tl.where(block < vocab_size, e, 0.0)
+ 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)
+
+
+@triton.jit
+def _ranks_kernel(
+ output_ptr,
+ logits_ptr,
+ logits_stride,
+ token_ids_ptr,
+ vocab_size,
+ BLOCK_SIZE: tl.constexpr,
+):
+ req_idx = tl.program_id(0)
+ row_ptr = logits_ptr + req_idx * logits_stride
+
+ token_id = tl.load(token_ids_ptr + req_idx)
+ x = tl.load(row_ptr + token_id)
+
+ n = 0
+ for i in range(0, vocab_size, BLOCK_SIZE):
+ block = i + tl.arange(0, BLOCK_SIZE)
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
+ n += tl.sum((logits > x).to(tl.int32))
+ tl.store(output_ptr + req_idx, n)
+
+
+def compute_token_logprobs(
+ logits: torch.Tensor,
+ token_ids: torch.Tensor,
+) -> torch.Tensor:
+ batch_size = logits.shape[0]
+ vocab_size = logits.shape[1]
+ token_ids = token_ids.to(torch.int64)
+ num_logprobs = token_ids.shape[1]
+ logprobs = torch.empty(
+ batch_size,
+ num_logprobs,
+ dtype=torch.float32,
+ device=logits.device,
+ )
+ _topk_log_softmax_kernel[(batch_size,)](
+ logprobs,
+ logits,
+ logits.stride(0),
+ token_ids,
+ num_logprobs,
+ vocab_size,
+ BLOCK_SIZE=1024, # type: ignore
+ PADDED_TOPK=triton.next_power_of_2(num_logprobs),
+ )
+ return logprobs
+
+
+def compute_topk_logprobs(
+ logits: torch.Tensor,
+ num_logprobs: int,
+ sampled_token_ids: torch.Tensor,
+) -> LogprobsTensors:
+ assert num_logprobs >= 0
+ batch_size, vocab_size = logits.shape
+ if num_logprobs == 0:
+ logprob_token_ids = sampled_token_ids.unsqueeze(-1)
+ else:
+ topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
+ logprob_token_ids = torch.cat(
+ (sampled_token_ids.unsqueeze(-1), topk_indices), dim=1
+ )
+
+ # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full
+ # logprobs tensor. Instead, we only compute and return the logprobs of
+ # the topk + 1 tokens.
+ logprobs = compute_token_logprobs(logits, logprob_token_ids)
+ token_ranks = torch.empty(
+ batch_size,
+ dtype=torch.int64,
+ device=logits.device,
+ )
+ _ranks_kernel[(batch_size,)](
+ token_ranks,
+ logits,
+ logits.stride(0),
+ sampled_token_ids,
+ vocab_size,
+ BLOCK_SIZE=8192, # type: ignore
+ )
+ return LogprobsTensors(
+ logprob_token_ids=logprob_token_ids,
+ logprobs=logprobs,
+ selected_token_ranks=token_ranks,
+ )
+
+
+def compute_prompt_logprobs(
+ prompt_token_ids: torch.Tensor,
+ prompt_hidden_states: torch.Tensor,
+ logits_fn: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+ # Since materializing the full prompt logits can take too much memory,
+ # we compute it in chunks.
+ CHUNK_SIZE = 1024
+ logprobs = []
+ ranks = []
+ prompt_token_ids = prompt_token_ids.to(torch.int64)
+ for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
+ end_idx = start_idx + CHUNK_SIZE
+ # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
+ prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
+ prompt_logprobs = compute_topk_logprobs(
+ prompt_logits,
+ 0, # num_logprobs
+ prompt_token_ids[start_idx:end_idx],
+ )
+ logprobs.append(prompt_logprobs.logprobs)
+ ranks.append(prompt_logprobs.selected_token_ranks)
+
+ logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
+ ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
+ return logprobs, ranks
diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sample.py b/vllm/v1/worker/gpu/spec_decode/rejection_sample.py
new file mode 100644
index 0000000000000..8a7bf28bacbd4
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sample.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _rejection_sample_kernel(
+ sampled_ptr, # [num_reqs, num_speculative_steps + 1]
+ sampled_stride,
+ num_sampled_ptr, # [num_reqs]
+ target_sampled_ptr, # [num_draft_tokens + num_reqs]
+ input_ids_ptr, # [num_draft_tokens + num_reqs]
+ cu_num_logits_ptr, # [num_reqs + 1]
+):
+ req_idx = tl.program_id(0)
+ start_idx = tl.load(cu_num_logits_ptr + req_idx)
+ end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
+ num_tokens = end_idx - start_idx
+
+ num_sampled = 0
+ rejected = False
+ for i in range(num_tokens - 1):
+ if not rejected:
+ target_sampled = tl.load(target_sampled_ptr + start_idx + i)
+ draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1)
+ tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled)
+ num_sampled += 1
+ if target_sampled != draft_sampled:
+ rejected = True
+ if not rejected:
+ target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
+ tl.store(
+ sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled
+ )
+ num_sampled += 1
+ tl.store(num_sampled_ptr + req_idx, num_sampled)
+
+
+def rejection_sample(
+ # [num_draft_tokens + num_reqs]
+ target_sampled: torch.Tensor,
+ # [num_draft_tokens + num_reqs]
+ input_ids: torch.Tensor,
+ # [num_reqs + 1]
+ cu_num_logits: torch.Tensor,
+ num_speculative_steps: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ num_reqs = cu_num_logits.shape[0] - 1
+ sampled = torch.empty(
+ num_reqs,
+ num_speculative_steps + 1,
+ dtype=target_sampled.dtype,
+ device=target_sampled.device,
+ )
+ num_sampled = torch.empty(
+ num_reqs,
+ dtype=torch.int32,
+ device=target_sampled.device,
+ )
+ _rejection_sample_kernel[(num_reqs,)](
+ sampled,
+ sampled.stride(0),
+ num_sampled,
+ target_sampled,
+ input_ids,
+ cu_num_logits,
+ num_warps=1,
+ )
+ return sampled, num_sampled
diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py
new file mode 100644
index 0000000000000..513d45d95d7cd
--- /dev/null
+++ b/vllm/v1/worker/gpu/states.py
@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import SamplingParams
+from vllm.triton_utils import tl, triton
+from vllm.v1.outputs import LogprobsTensors
+from vllm.v1.utils import CpuGpuBuffer
+
+_NP_INT64_MIN = np.iinfo(np.int64).min
+_NP_INT64_MAX = np.iinfo(np.int64).max
+NO_LORA_ID = 0
+
+
+@dataclass
+class SamplingMetadata:
+ temperature: torch.Tensor
+
+ top_p: torch.Tensor | None
+ top_k: torch.Tensor | None
+
+ seeds: torch.Tensor
+ pos: torch.Tensor
+
+ # None means no logprobs, 0 means sampled token logprobs only
+ max_num_logprobs: int | None
+
+ @classmethod
+ def make_dummy(
+ cls,
+ num_reqs: int,
+ device: torch.device,
+ ) -> "SamplingMetadata":
+ assert num_reqs > 0
+ temperature = torch.zeros(num_reqs, dtype=torch.float32, device=device)
+ temperature[0] = 0.5
+ # TODO(woosuk): Use top-p and top-k for dummy sampler.
+ # Currently, they are disabled because of memory usage.
+ # top_p = torch.full((num_reqs,), 0.95, dtype=torch.float32, device=device)
+ # top_k = torch.full((num_reqs,), 20, dtype=torch.int32, device=device)
+ top_p = None
+ top_k = None
+ seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
+ pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
+ max_num_logprobs = 20
+
+ return cls(
+ temperature=temperature,
+ top_p=top_p,
+ top_k=top_k,
+ seeds=seeds,
+ pos=pos,
+ max_num_logprobs=max_num_logprobs,
+ )
+
+
+class RequestState:
+ def __init__(
+ self,
+ max_num_reqs: int,
+ max_model_len: int,
+ max_num_batched_tokens: int,
+ num_speculative_steps: int,
+ vocab_size: int,
+ device: torch.device,
+ pin_memory: bool,
+ ):
+ self.max_num_reqs = max_num_reqs
+ self.max_model_len = max_model_len
+ self.max_num_batched_tokens = max_num_batched_tokens
+ self.num_speculative_steps = num_speculative_steps
+ self.vocab_size = vocab_size
+ self.device = device
+ self.pin_memory = pin_memory
+
+ self.req_id_to_index: dict[str, int] = {}
+ self.index_to_req_id: dict[int, str] = {}
+ self.free_indices = list(range(max_num_reqs))
+ self.extra_data: dict[str, ExtraData] = {}
+
+ self.prompt_len = np.zeros(self.max_num_reqs, dtype=np.int32)
+ self.prefill_token_ids = np.zeros(
+ (self.max_num_reqs, self.max_model_len),
+ dtype=np.int32,
+ )
+ self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
+
+ # Number of computed tokens.
+ self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
+ self.num_computed_tokens = torch.zeros(
+ self.max_num_reqs, dtype=torch.int32, device=device
+ )
+
+ # Last sampled tokens.
+ self.last_sampled_tokens = torch.zeros(
+ self.max_num_reqs,
+ 1,
+ dtype=torch.int64,
+ device=device,
+ )
+
+ # Draft tokens.
+ self.draft_tokens = torch.zeros(
+ self.max_num_reqs,
+ self.num_speculative_steps,
+ dtype=torch.int64,
+ device=device,
+ )
+
+ # LoRA.
+ self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
+ self.lora_ids.fill(NO_LORA_ID)
+
+ # Sampling parameters.
+ self.temperature = self._make_param(self.max_num_reqs, torch.float32)
+ self.top_p = self._make_param(self.max_num_reqs, torch.float32)
+ self.top_k = self._make_param(self.max_num_reqs, torch.int32)
+ self.seeds = self._make_param(self.max_num_reqs, torch.int64)
+
+ self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
+ # -1 means no logprobs are requested.
+ self.num_logprobs.fill(-1)
+ self.needs_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
+
+ def _make_param(self, size: int, dtype: torch.dtype) -> "Param":
+ return Param(size, dtype=dtype, device=self.device, pin_memory=self.pin_memory)
+
+ def _make_buffer(self, size: int, dtype: torch.dtype) -> CpuGpuBuffer:
+ return CpuGpuBuffer(
+ size, dtype=dtype, device=self.device, pin_memory=self.pin_memory
+ )
+
+ @property
+ def num_reqs(self) -> int:
+ return len(self.req_id_to_index)
+
+ def add_request(
+ self,
+ req_id: str,
+ prompt_len: int,
+ prefill_token_ids: list[int],
+ num_computed_tokens: int,
+ sampling_params: SamplingParams,
+ lora_request: LoRARequest | None,
+ ) -> None:
+ assert len(self.free_indices) > 0, "No free indices"
+ req_idx = self.free_indices.pop()
+ self.req_id_to_index[req_id] = req_idx
+ self.index_to_req_id[req_idx] = req_id
+ self.extra_data[req_id] = ExtraData(lora_request)
+
+ self.prompt_len[req_idx] = prompt_len
+ prefill_len = len(prefill_token_ids)
+ assert prefill_len >= prompt_len, (
+ f"prefill_len {prefill_len} < prompt_len {prompt_len}"
+ )
+ self.prefill_len.np[req_idx] = prefill_len
+ self.prefill_token_ids[req_idx, :prefill_len] = prefill_token_ids
+
+ self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
+ # FIXME(woosuk): This triggers a GPU operation whenever adding a new request.
+ # Optimize this.
+ self.num_computed_tokens[req_idx] = num_computed_tokens
+
+ if lora_request is not None:
+ self.lora_ids[req_idx] = lora_request.lora_int_id
+ else:
+ self.lora_ids[req_idx] = NO_LORA_ID
+
+ self.temperature.np[req_idx] = sampling_params.temperature
+ self.top_p.np[req_idx] = sampling_params.top_p
+ if 0 < sampling_params.top_k < self.vocab_size:
+ top_k = sampling_params.top_k
+ else:
+ top_k = self.vocab_size
+ self.top_k.np[req_idx] = top_k
+
+ if sampling_params.seed is not None:
+ seed = sampling_params.seed
+ else:
+ seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)
+ self.seeds.np[req_idx] = seed
+
+ if sampling_params.logprobs is not None:
+ num_logprobs = sampling_params.logprobs
+ else:
+ num_logprobs = -1
+ self.num_logprobs[req_idx] = num_logprobs
+
+ # For now, only support prompt logprobs for the prompt tokens.
+ needs_prompt_logprobs = sampling_params.prompt_logprobs is not None
+ self.needs_prompt_logprobs[req_idx] = needs_prompt_logprobs
+
+ def remove_request(self, req_id: str) -> None:
+ self.extra_data.pop(req_id, None)
+ req_idx = self.req_id_to_index.pop(req_id, None)
+ if req_idx is None:
+ # Request not found.
+ return
+ self.index_to_req_id.pop(req_idx, None)
+ self.free_indices.append(req_idx)
+
+ def make_sampling_metadata(
+ self,
+ idx_mapping: np.ndarray,
+ pos: torch.Tensor,
+ ) -> SamplingMetadata:
+ temperature = self.temperature.np[idx_mapping]
+ temperature = self.temperature.copy_np_to_gpu(temperature)
+
+ top_p = self.top_p.np[idx_mapping]
+ no_top_p = np.all(top_p == 1.0)
+ top_p = self.top_p.copy_np_to_gpu(top_p) if not no_top_p else None
+
+ top_k = self.top_k.np[idx_mapping]
+ no_top_k = np.all(top_k == self.vocab_size)
+ top_k = self.top_k.copy_np_to_gpu(top_k) if not no_top_k else None
+
+ seeds = self.seeds.np[idx_mapping]
+ seeds = self.seeds.copy_np_to_gpu(seeds)
+
+ num_logprobs = self.num_logprobs[idx_mapping]
+ max_num_logprobs: int | None = int(np.max(num_logprobs))
+ if max_num_logprobs == -1:
+ max_num_logprobs = None
+
+ return SamplingMetadata(
+ temperature=temperature,
+ top_p=top_p,
+ top_k=top_k,
+ seeds=seeds,
+ pos=pos,
+ max_num_logprobs=max_num_logprobs,
+ )
+
+ def expand_sampling_metadata(
+ self,
+ sampling_metadata: SamplingMetadata,
+ cu_num_logits: torch.Tensor,
+ ) -> SamplingMetadata:
+ # For draft tokens, we need to expand the sampling param tensors as
+ # each request samples multiple tokens in each step.
+ return expand_sampling_metadata(
+ sampling_metadata, cu_num_logits, self.num_speculative_steps
+ )
+
+ def make_lora_inputs(
+ self,
+ req_ids: list[str],
+ idx_mapping: np.ndarray,
+ num_scheduled_tokens: np.ndarray,
+ ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
+ lora_ids = self.lora_ids[idx_mapping]
+ prompt_lora_mapping = tuple(lora_ids)
+ token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
+
+ active_lora_requests: set[LoRARequest] = set()
+ for req_id in req_ids:
+ lora_request = self.extra_data[req_id].lora_request
+ if lora_request is not None:
+ active_lora_requests.add(lora_request)
+ return prompt_lora_mapping, token_lora_mapping, active_lora_requests
+
+
+class Param:
+ def __init__(
+ self,
+ size: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ pin_memory: bool,
+ ):
+ self.buffer = CpuGpuBuffer(
+ size,
+ dtype=dtype,
+ device=device,
+ pin_memory=pin_memory,
+ )
+ self.np = np.zeros_like(self.buffer.np)
+
+ def copy_np_to_gpu(self, x: np.ndarray) -> torch.Tensor:
+ n = x.shape[0]
+ self.buffer.np[:n] = x
+ return self.buffer.copy_to_gpu(n)
+
+
+@dataclass
+class ExtraData:
+ lora_request: LoRARequest | None
+ in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)
+
+
+# NOTE(woosuk): Re-compilation can happen at runtime since top_p and top_k can be None.
+@triton.jit
+def _expand_sampling_metadata_kernel(
+ temp_ptr,
+ expanded_temp_ptr,
+ top_p_ptr,
+ expanded_top_p_ptr,
+ top_k_ptr,
+ expanded_top_k_ptr,
+ seeds_ptr,
+ expanded_seeds_ptr,
+ cu_num_logits_ptr,
+ BLOCK_SIZE: tl.constexpr,
+):
+ req_idx = tl.program_id(0)
+ start_idx = tl.load(cu_num_logits_ptr + req_idx)
+ end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
+ num_tokens = end_idx - start_idx
+
+ block = tl.arange(0, BLOCK_SIZE)
+ mask = block < num_tokens
+
+ temp = tl.load(temp_ptr + req_idx)
+ tl.store(expanded_temp_ptr + start_idx + block, temp, mask=mask)
+
+ if top_p_ptr is not None:
+ top_p = tl.load(top_p_ptr + req_idx)
+ tl.store(expanded_top_p_ptr + start_idx + block, top_p, mask=mask)
+
+ if top_k_ptr is not None:
+ top_k = tl.load(top_k_ptr + req_idx)
+ tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)
+
+ seed = tl.load(seeds_ptr + req_idx)
+ tl.store(expanded_seeds_ptr + start_idx + block, seed, mask=mask)
+
+
+def expand_sampling_metadata(
+ sampling_metadata: SamplingMetadata,
+ cu_num_logits: torch.Tensor,
+ num_speculative_steps: int,
+) -> SamplingMetadata:
+ total_num_logits = sampling_metadata.pos.shape[0]
+ create_empty = lambda x: x.new_empty(total_num_logits) if x is not None else None
+ expanded_temp = create_empty(sampling_metadata.temperature)
+ expanded_top_p = create_empty(sampling_metadata.top_p)
+ expanded_top_k = create_empty(sampling_metadata.top_k)
+ expanded_seeds = create_empty(sampling_metadata.seeds)
+
+ num_reqs = cu_num_logits.shape[0] - 1
+ _expand_sampling_metadata_kernel[(num_reqs,)](
+ sampling_metadata.temperature,
+ expanded_temp,
+ sampling_metadata.top_p,
+ expanded_top_p,
+ sampling_metadata.top_k,
+ expanded_top_k,
+ sampling_metadata.seeds,
+ expanded_seeds,
+ cu_num_logits,
+ BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
+ )
+ return SamplingMetadata(
+ temperature=expanded_temp,
+ top_p=expanded_top_p,
+ top_k=expanded_top_k,
+ seeds=expanded_seeds,
+ pos=sampling_metadata.pos,
+ max_num_logprobs=sampling_metadata.max_num_logprobs,
+ )
diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py
new file mode 100644
index 0000000000000..83051b0ed33ff
--- /dev/null
+++ b/vllm/v1/worker/gpu/structured_outputs.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import numpy as np
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.input_batch import InputBuffers
+
+
+def apply_grammar_bitmask(
+ logits: torch.Tensor,
+ req_ids: list[str],
+ grammar_req_ids: list[str],
+ grammar_bitmask: np.ndarray,
+ input_buffers: InputBuffers,
+) -> None:
+ input_buffers.grammar_bitmask.np[: grammar_bitmask.shape[0]] = grammar_bitmask
+ input_buffers.grammar_bitmask.copy_to_gpu(grammar_bitmask.shape[0])
+
+ batch_size = logits.shape[0]
+ grammar_req_id_to_idx = {req_id: i for i, req_id in enumerate(grammar_req_ids)}
+ # logits -> bitmask mapping
+ mapping = [grammar_req_id_to_idx.get(req_id, -1) for req_id in req_ids]
+ input_buffers.bitmask_indices.np[:batch_size] = mapping
+ input_buffers.bitmask_indices.copy_to_gpu(batch_size)
+
+ vocab_size = logits.shape[-1]
+ BLOCK_SIZE = 8192
+ grid = (batch_size, triton.cdiv(vocab_size, BLOCK_SIZE))
+ _apply_grammar_bitmask_kernel[grid](
+ logits,
+ logits.stride(0),
+ input_buffers.grammar_bitmask.gpu,
+ input_buffers.grammar_bitmask.gpu.stride(0),
+ input_buffers.bitmask_indices.gpu,
+ vocab_size,
+ BLOCK_SIZE=BLOCK_SIZE,
+ )
+
+
+# Adapted from
+# https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py
+@triton.jit
+def _apply_grammar_bitmask_kernel(
+ logits_ptr,
+ logits_stride,
+ bitmask_ptr,
+ bitmask_stride,
+ bitmask_indices_ptr,
+ vocab_size,
+ BLOCK_SIZE: tl.constexpr,
+):
+ logits_idx = tl.program_id(0)
+ bitmask_idx = tl.load(bitmask_indices_ptr + logits_idx)
+ if bitmask_idx == -1:
+ # No bitmask to apply.
+ return
+
+ # Load the bitmask.
+ block_id = tl.program_id(1)
+ bitmask_offset = (block_id * BLOCK_SIZE) // 32 + tl.arange(0, BLOCK_SIZE // 32)
+ packed_bitmask = tl.load(
+ bitmask_ptr + bitmask_idx * bitmask_stride + bitmask_offset,
+ mask=bitmask_offset < bitmask_stride,
+ )
+ # Unpack the bitmask.
+ bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0
+ bitmask = bitmask.reshape(BLOCK_SIZE)
+
+ # Apply the bitmask to the logits.
+ block_offset = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ tl.store(
+ logits_ptr + logits_idx * logits_stride + block_offset,
+ -float("inf"),
+ mask=bitmask & (block_offset < vocab_size),
+ )
diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py
index 7cf6afa3fc371..d6fef450c028a 100644
--- a/vllm/v1/worker/gpu_input_batch.py
+++ b/vllm/v1/worker/gpu_input_batch.py
@@ -87,7 +87,7 @@ class InputBatch:
is_spec_decode: bool = False,
is_pooling_model: bool = False,
num_speculative_tokens: int = 0,
- dcp_kv_cache_interleave_size: int = 1,
+ cp_kv_cache_interleave_size: int = 1,
):
self.is_pooling_model = is_pooling_model
self.is_spec_decode = is_spec_decode
@@ -141,7 +141,7 @@ class InputBatch:
block_sizes=block_sizes,
kernel_block_sizes=kernel_block_sizes,
num_speculative_tokens=num_speculative_tokens,
- dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
)
# Sampling-related.
@@ -219,9 +219,6 @@ class InputBatch:
self.generators: dict[int, torch.Generator] = {}
self.num_logprobs: dict[str, int] = {}
- # NOTE(rob): num_prompt_logprobs only includes reqs
- # that are currently in the prefill phase.
- self.num_prompt_logprobs: dict[str, int] = {}
# To accumulate prompt logprobs tensor chunks across prefill steps.
self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
@@ -251,7 +248,7 @@ class InputBatch:
self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids
# Store last speculative tokens for sampler.
- self.spec_token_ids: list[list[int] | None] = []
+ self.spec_token_ids: list[list[int]] = [[] for _ in range(max_num_reqs)]
# This is updated each time the batch constituents change.
self.sampling_metadata = self._make_sampling_metadata()
@@ -265,7 +262,7 @@ class InputBatch:
# ids from prior step, if required by current sampling params
# (e.g. penalties).
self.sampled_token_ids_cpu: torch.Tensor | None = None
- self.async_copy_ready_event: torch.cuda.Event | None = None
+ self.async_copy_ready_event: torch.Event | None = None
@property
def req_ids(self) -> list[str]:
@@ -313,7 +310,7 @@ class InputBatch:
else:
self._req_ids[req_index] = req_id
self.req_output_token_ids[req_index] = request.output_token_ids
- self.spec_token_ids[req_index] = []
+ self.spec_token_ids[req_index].clear()
self.req_id_to_index[req_id] = req_index
@@ -385,12 +382,6 @@ class InputBatch:
if sampling_params.logprobs == -1
else sampling_params.logprobs
)
- if sampling_params.prompt_logprobs is not None:
- self.num_prompt_logprobs[req_id] = (
- self.vocab_size
- if sampling_params.prompt_logprobs == -1
- else sampling_params.prompt_logprobs
- )
if sampling_params.allowed_token_ids:
self.has_allowed_token_ids.add(req_id)
@@ -462,7 +453,7 @@ class InputBatch:
self.batch_update_builder.removed_append(req_index)
self._req_ids[req_index] = None
self.req_output_token_ids[req_index] = None
- self.spec_token_ids[req_index] = None
+ self.spec_token_ids[req_index].clear()
# LoRA
lora_id = self.request_lora_mapping[req_index]
@@ -488,7 +479,6 @@ class InputBatch:
self.repetition_penalties_reqs.discard(req_id)
self.generators.pop(req_index, None)
self.num_logprobs.pop(req_id, None)
- self.num_prompt_logprobs.pop(req_id, None)
self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
self.has_allowed_token_ids.discard(req_id)
@@ -654,9 +644,15 @@ class InputBatch:
self.req_output_token_ids[last_req_index] = None
self.req_id_to_index[req_id] = empty_index
- spec_token_ids = self.spec_token_ids[last_req_index]
- self.spec_token_ids[empty_index] = spec_token_ids
- self.spec_token_ids[last_req_index] = None
+ if last_req_index != empty_index:
+ (
+ self.spec_token_ids[last_req_index],
+ self.spec_token_ids[empty_index],
+ ) = (
+ self.spec_token_ids[empty_index],
+ self.spec_token_ids[last_req_index],
+ )
+ self.spec_token_ids[last_req_index].clear()
num_tokens = self.num_tokens[last_req_index]
self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
@@ -891,7 +887,7 @@ class InputBatch:
def set_async_sampled_token_ids(
self,
sampled_token_ids_cpu: torch.Tensor,
- async_copy_ready_event: torch.cuda.Event,
+ async_copy_ready_event: torch.Event,
) -> None:
"""
In async scheduling case, store ref to sampled_token_ids_cpu
@@ -966,10 +962,6 @@ class InputBatch:
def max_num_logprobs(self) -> int | None:
return max(self.num_logprobs.values()) if self.num_logprobs else None
- @property
- def no_prompt_logprob(self) -> bool:
- return not self.num_prompt_logprobs
-
@property
def no_allowed_token_ids(self) -> bool:
return len(self.has_allowed_token_ids) == 0
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index 0102ca4739ad5..6a54e02f861e9 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -5,7 +5,7 @@ import gc
import itertools
import time
from collections import defaultdict
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from copy import copy, deepcopy
from functools import reduce
@@ -53,6 +53,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
from vllm.model_executor.models.interfaces import (
+ SupportsMRoPE,
SupportsMultiModal,
is_mixture_of_experts,
supports_eagle3,
@@ -126,6 +127,7 @@ from vllm.v1.outputs import (
)
from vllm.v1.pool.metadata import PoolingMetadata
from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs
+from vllm.v1.sample.logits_processor.interface import LogitsProcessor
from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.sample.rejection_sampler import RejectionSampler
from vllm.v1.sample.sampler import Sampler
@@ -185,7 +187,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
self._invalid_req_indices = invalid_req_indices
# Event on the copy stream so we can synchronize the non-blocking copy.
- self.async_copy_ready_event = torch.cuda.Event()
+ self.async_copy_ready_event = torch.Event()
# Keep a reference to the device tensor to avoid it being
# deallocated until we finish copying it to the host.
@@ -219,16 +221,14 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
del self._sampled_token_ids
max_gen_len = self.sampled_token_ids_cpu.shape[-1]
if max_gen_len == 1:
- valid_sampled_token_ids: list[np.ndarray] = [
- row for row in self.sampled_token_ids_cpu.numpy()
- ]
+ valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()
else:
valid_sampled_token_ids = RejectionSampler.parse_output(
self.sampled_token_ids_cpu,
self.vocab_size,
)
for i in self._invalid_req_indices:
- valid_sampled_token_ids[i] = np.array([])
+ valid_sampled_token_ids[i].clear()
output = self._model_runner_output
output.sampled_token_ids = valid_sampled_token_ids
@@ -349,6 +349,9 @@ class GPUModelRunner(
# self.model: nn.Module # Set after load_model
# Initialize in initialize_kv_cache
self.kv_caches: list[torch.Tensor] = []
+ # Initialize in initialize_kv_cache_tensors
+ self.cross_layers_kv_cache: torch.Tensor | None = None
+ self.cross_layers_attn_backend: type[AttentionBackend] | None = None
# indexes: [kv_cache_group_id][attn_group]
self.attn_groups: list[list[AttentionGroup]] = []
# self.kv_cache_config: KVCacheConfig
@@ -390,6 +393,9 @@ class GPUModelRunner(
# Request states.
self.requests: dict[str, CachedRequestState] = {}
+ # NOTE(rob): num_prompt_logprobs only includes reqs
+ # that are currently in the prefill phase.
+ self.num_prompt_logprobs: dict[str, int] = {}
self.comm_stream = torch.cuda.Stream()
# Input Batch
@@ -401,7 +407,10 @@ class GPUModelRunner(
# solution, we initialize the input batch here, and re-initialize it
# in `initialize_kv_cache` if the block_sizes here is different from
# the block_sizes in the kv cache config.
- custom_logitsprocs = model_config.logits_processors
+ logits_processors = model_config.logits_processors
+ custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (
+ tuple(logits_processors) if logits_processors is not None else ()
+ )
self.input_batch = InputBatch(
max_num_reqs=self.max_num_reqs,
# We need to use the encoder length for encoder-decoer
@@ -425,7 +434,7 @@ class GPUModelRunner(
# uses output token ids so we set this conservatively.
logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
is_pooling_model=self.is_pooling_model,
- dcp_kv_cache_interleave_size=self.parallel_config.dcp_kv_cache_interleave_size,
+ cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
)
self.use_async_scheduling = self.scheduler_config.async_scheduling
@@ -434,10 +443,10 @@ class GPUModelRunner(
self.async_output_copy_stream: torch.cuda.Stream | None = None
# cuda event to synchronize use of reused CPU tensors between steps
# when async scheduling is enabled.
- self.prepare_inputs_event: torch.cuda.Event | None = None
+ self.prepare_inputs_event: torch.Event | None = None
if self.use_async_scheduling:
self.async_output_copy_stream = torch.cuda.Stream()
- self.prepare_inputs_event = torch.cuda.Event()
+ self.prepare_inputs_event = torch.Event()
# self.cudagraph_batch_sizes sorts in ascending order.
if (
@@ -548,7 +557,7 @@ class GPUModelRunner(
# Cached outputs.
self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
- self.transfer_event = torch.cuda.Event()
+ self.transfer_event = torch.Event()
self.sampled_token_ids_pinned_cpu = torch.empty(
(self.max_num_reqs, 1),
dtype=torch.int64,
@@ -558,10 +567,10 @@ class GPUModelRunner(
# Pre-allocated tensor for copying valid sampled token counts to CPU,
# with dedicated stream for overlapping and event for coordination.
- self.valid_sampled_token_count_event: torch.cuda.Event | None = None
+ self.valid_sampled_token_count_event: torch.Event | None = None
self.valid_sampled_token_count_copy_stream: torch.cuda.Stream | None = None
if self.use_async_scheduling and self.num_spec_tokens:
- self.valid_sampled_token_count_event = torch.cuda.Event()
+ self.valid_sampled_token_count_event = torch.Event()
self.valid_sampled_token_count_copy_stream = torch.cuda.Stream()
self.valid_sampled_token_count_cpu = torch.empty(
self.max_num_reqs,
@@ -681,6 +690,7 @@ class GPUModelRunner(
# Remove finished requests from the cached states.
for req_id in scheduler_output.finished_req_ids:
self.requests.pop(req_id, None)
+ self.num_prompt_logprobs.pop(req_id, None)
# Remove the finished requests from the persistent batch.
# NOTE(woosuk): There could be an edge case where finished_req_ids and
# scheduled_req_ids overlap. This happens when a request is aborted and
@@ -749,6 +759,13 @@ class GPUModelRunner(
)
self.requests[req_id] = req_state
+ if sampling_params and sampling_params.prompt_logprobs is not None:
+ self.num_prompt_logprobs[req_id] = (
+ self.input_batch.vocab_size
+ if sampling_params.prompt_logprobs == -1
+ else sampling_params.prompt_logprobs
+ )
+
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
if self.uses_mrope:
self._init_mrope_positions(req_state)
@@ -891,7 +908,8 @@ class GPUModelRunner(
# conform to the schema. This can result in
# scheduler_output.scheduled_spec_decode_tokens being empty,
# even when speculative decoding is enabled.
- self.input_batch.spec_token_ids[req_index] = spec_token_ids
+ self.input_batch.spec_token_ids[req_index].clear()
+ self.input_batch.spec_token_ids[req_index].extend(spec_token_ids)
# there are no draft tokens with async scheduling,
# we clear the spec_decoding info in scheduler_output and
@@ -955,9 +973,13 @@ class GPUModelRunner(
def _init_mrope_positions(self, req_state: CachedRequestState):
model = self.get_model()
assert supports_mrope(model), "M-RoPE support is not implemented."
+ assert req_state.prompt_token_ids is not None, (
+ "M-RoPE requires prompt_token_ids to be available."
+ )
+ mrope_model = cast(SupportsMRoPE, model)
req_state.mrope_positions, req_state.mrope_position_delta = (
- model.get_mrope_input_positions(
+ mrope_model.get_mrope_input_positions(
req_state.prompt_token_ids,
req_state.mm_features,
)
@@ -1434,7 +1456,7 @@ class GPUModelRunner(
self.seq_lens.cpu[:num_reqs],
self.dcp_world_size,
self.dcp_rank,
- self.parallel_config.dcp_kv_cache_interleave_size,
+ self.parallel_config.cp_kv_cache_interleave_size,
)
self.dcp_local_seq_lens.copy_to_gpu(num_reqs)
@@ -1450,9 +1472,12 @@ class GPUModelRunner(
num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[
:num_reqs
]
- dcp_local_seq_lens = (
- self.dcp_local_seq_lens.gpu[:num_reqs] if self.dcp_world_size > 1 else None
- )
+
+ dcp_local_seq_lens, dcp_local_seq_lens_cpu = None, None
+ if self.dcp_world_size > 1:
+ dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs]
+ dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[:num_reqs]
+
spec_decode_common_attn_metadata = None
if for_cudagraph_capture:
@@ -1520,6 +1545,7 @@ class GPUModelRunner(
causal=True,
encoder_seq_lens=encoder_seq_lens,
dcp_local_seq_lens=dcp_local_seq_lens,
+ dcp_local_seq_lens_cpu=dcp_local_seq_lens_cpu,
)
if self.speculative_config and spec_decode_common_attn_metadata is None:
@@ -1754,6 +1780,7 @@ class GPUModelRunner(
dst_start = mrope_pos_ptr
dst_end = mrope_pos_ptr + completion_part_len
+ assert req.mrope_position_delta is not None
MRotaryEmbedding.get_next_input_positions_tensor(
out=self.mrope_positions.np,
out_offset=dst_start,
@@ -1899,20 +1926,24 @@ class GPUModelRunner(
for mm_input_id in encoder_input_ids:
mm_feature = req_state.mm_features[mm_input_id]
+ if mm_feature.data is None:
+ continue
mm_hash = mm_feature.identifier
mm_kwargs.append(mm_feature.data)
mm_hashes_pos.append((mm_hash, mm_feature.mm_position))
return mm_kwargs, mm_hashes_pos
- def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
+ def _execute_mm_encoder(
+ self, scheduler_output: "SchedulerOutput"
+ ) -> list[torch.Tensor]:
# Batch the multi-modal inputs using the helper method.
mm_kwargs, mm_hashes_pos = self._batch_mm_kwargs_from_scheduler(
scheduler_output
)
if not mm_kwargs:
- return
+ return []
# Batch mm inputs as much as we can: if a request in the batch has
# multiple modalities or a different modality than the previous one,
@@ -1922,7 +1953,7 @@ class GPUModelRunner(
# multimodal inputs. The proper solution should be reordering the
# encoder outputs.
model = cast(SupportsMultiModal, self.model)
- encoder_outputs = []
+ encoder_outputs: list[torch.Tensor] = []
for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
mm_kwargs,
device=self.device,
@@ -1930,7 +1961,7 @@ class GPUModelRunner(
merge_by_field_config=model.merge_by_field_config,
multimodal_cpu_fields=model.multimodal_cpu_fields,
):
- curr_group_outputs = []
+ curr_group_outputs: list[torch.Tensor] = []
# EVS-related change.
# (ekhvedchenia): Temporary hack to limit peak memory usage when
@@ -1972,7 +2003,7 @@ class GPUModelRunner(
# 2. A list or tuple (length: num_items) of tensors,
# each of shape (feature_size, hidden_size) in case the feature
# size is dynamic depending on the input multimodal items.
- curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)
+ curr_group_outputs = model.embed_multimodal(**mm_kwargs_group) # type: ignore[assignment]
sanity_check_mm_encoder_outputs(
curr_group_outputs,
@@ -1989,6 +2020,8 @@ class GPUModelRunner(
logger.debug("Finish execute for mm hash %s", mm_hash)
self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash)
+ return encoder_outputs
+
def _gather_mm_embeddings(
self,
scheduler_output: "SchedulerOutput",
@@ -2077,38 +2110,6 @@ class GPUModelRunner(
return mm_embeds, is_mm_embed
- def _extract_encoder_inputs(
- self,
- scheduler_output: "SchedulerOutput",
- ) -> dict[str, torch.Tensor]:
- """Extract encoder inputs for encoder-decoder models.
-
- This method extracts multimodal input features from scheduled encoder
- inputs and formats them for the encoder-decoder model forward pass.
- """
- # Batch the multi-modal inputs using the helper method.
- mm_kwargs, _ = self._batch_mm_kwargs_from_scheduler(scheduler_output)
-
- if not mm_kwargs:
- return {}
-
- # Group MM kwargs by modality and extract features
- model = cast(SupportsMultiModal, self.model)
- encoder_features = {}
- for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
- mm_kwargs,
- device=self.device,
- pin_memory=self.pin_memory,
- merge_by_field_config=model.merge_by_field_config,
- multimodal_cpu_fields=model.multimodal_cpu_fields,
- ):
- # Add the grouped features to encoder_features dict
- # This allows the model to receive them as kwargs (e.g.,
- # input_features=...)
- encoder_features.update(mm_kwargs_group)
-
- return encoder_features
-
def get_model(self) -> nn.Module:
# get raw model out of the cudagraph wrapper.
if isinstance(self.model, (CUDAGraphWrapper, UBatchWrapper)):
@@ -2172,7 +2173,7 @@ class GPUModelRunner(
def sync_and_slice_intermediate_tensors(
self,
num_tokens: int,
- intermediate_tensors: IntermediateTensors,
+ intermediate_tensors: IntermediateTensors | None,
sync_self: bool,
) -> IntermediateTensors:
assert self.intermediate_tensors is not None
@@ -2389,6 +2390,7 @@ class GPUModelRunner(
if is_first_rank:
intermediate_tensors = None
else:
+ assert intermediate_tensors is not None
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
num_input_tokens, intermediate_tensors, True
)
@@ -2397,8 +2399,13 @@ class GPUModelRunner(
self.model_config.is_encoder_decoder
and scheduler_output.scheduled_encoder_inputs
):
- encoder_inputs = self._extract_encoder_inputs(scheduler_output)
- model_kwargs.update(encoder_inputs)
+ # Run the encoder, just like we do with other multimodal inputs.
+ # For an encoder-decoder model, our processing here is a bit
+ # simpler, because the outputs are just passed to the decoder.
+ # We are not doing any prompt replacement. We also will only
+ # ever have a single encoder input.
+ encoder_outputs = self._execute_mm_encoder(scheduler_output)
+ model_kwargs.update({"encoder_outputs": encoder_outputs})
return (
input_ids,
@@ -2445,7 +2452,7 @@ class GPUModelRunner(
) -> tuple[
dict[str, int],
LogprobsLists | None,
- list[np.ndarray],
+ list[list[int]],
dict[str, LogprobsTensors | None],
list[str],
dict[str, int],
@@ -2470,8 +2477,9 @@ class GPUModelRunner(
num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]
sampled_token_ids = sampler_output.sampled_token_ids
+ logprobs_tensors = sampler_output.logprobs_tensors
invalid_req_indices = []
- valid_sampled_token_ids: list[np.ndarray]
+ cu_num_new_tokens: list[int] | None = None
if not self.use_async_scheduling:
# Get the valid generated tokens.
max_gen_len = sampled_token_ids.shape[-1]
@@ -2484,9 +2492,15 @@ class GPUModelRunner(
sampled_token_ids,
self.input_batch.vocab_size,
)
+ if logprobs_tensors:
+ # Needed for extracting logprobs when spec decoding.
+ # This must be done prior to discarding sampled tokens.
+ cu_num_new_tokens = [0]
+ for toks in valid_sampled_token_ids:
+ cu_num_new_tokens.append(cu_num_new_tokens[-1] + len(toks))
# Mask out the sampled tokens that should not be sampled.
for i in discard_sampled_tokens_req_indices:
- valid_sampled_token_ids[int(i)] = np.array([])
+ valid_sampled_token_ids[int(i)].clear()
else:
valid_sampled_token_ids = []
invalid_req_indices = discard_sampled_tokens_req_indices.tolist()
@@ -2511,29 +2525,15 @@ class GPUModelRunner(
# the sampled tokens back, because there's no direct communication
# between the first-stage worker and the last-stage worker.
req_ids = self.input_batch.req_ids
- logprobs_tensors = sampler_output.logprobs_tensors
- cu_num_accepted_tokens = (
- [0] if spec_decode_metadata and logprobs_tensors else None
- )
for req_idx in range(num_sampled_tokens):
- sampled_ids: np.ndarray | None
if self.use_async_scheduling:
- sampled_ids = (
- np.array([-1]) if req_idx not in invalid_req_indices_set else None
- )
+ sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None
else:
sampled_ids = valid_sampled_token_ids[req_idx]
- num_sampled_ids: int = (
- sampled_ids.shape[0] if sampled_ids is not None else 0
- )
+ num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0
- if cu_num_accepted_tokens is not None:
- cu_num_accepted_tokens.append(
- cu_num_accepted_tokens[-1] + num_sampled_ids
- )
-
- if sampled_ids is None or num_sampled_ids == 0:
+ if not sampled_ids:
continue
start_idx = self.input_batch.num_tokens_no_spec[req_idx]
@@ -2554,7 +2554,7 @@ class GPUModelRunner(
req_state.output_token_ids.extend(sampled_ids)
logprobs_lists = (
- logprobs_tensors.tolists(cu_num_accepted_tokens)
+ logprobs_tensors.tolists(cu_num_new_tokens)
if not self.use_async_scheduling and logprobs_tensors is not None
else None
)
@@ -2682,7 +2682,7 @@ class GPUModelRunner(
scheduler_output, self.vllm_config
)
if self.cache_config.kv_sharing_fast_prefill:
- assert not self.input_batch.num_prompt_logprobs, (
+ assert not self.num_prompt_logprobs, (
"--kv-sharing-fast-prefill produces incorrect "
"logprobs for prompt tokens, tokens, please disable "
"it when the requests need prompt logprobs"
@@ -2757,14 +2757,14 @@ class GPUModelRunner(
uniform_decode = (
max_num_scheduled_tokens == self.uniform_decode_query_len
) and (num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)
- batch_descriptor = BatchDescriptor(
+ batch_desc = BatchDescriptor(
num_tokens=num_input_tokens,
uniform_decode=uniform_decode,
has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,
)
cudagraph_runtime_mode, batch_descriptor = (
self.cudagraph_dispatcher.dispatch(
- batch_descriptor,
+ batch_desc,
use_cascade_attn=cascade_attn_prefix_lens is not None,
)
)
@@ -2848,15 +2848,15 @@ class GPUModelRunner(
else:
logits = self.model.compute_logits(sample_hidden_states)
- model_output_broadcast_data = {}
+ model_output_broadcast_data: dict[str, Any] = {}
if logits is not None:
model_output_broadcast_data["logits"] = logits.contiguous()
- model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+ broadcasted = get_pp_group().broadcast_tensor_dict(
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
)
- assert model_output_broadcast_data is not None
- logits = model_output_broadcast_data["logits"]
+ assert broadcasted is not None
+ logits = broadcasted["logits"]
self.execute_model_state = ExecuteModelState(
scheduler_output,
@@ -2881,7 +2881,7 @@ class GPUModelRunner(
if self.execute_model_state is None:
# Nothing to do (PP non-final rank case), output isn't used.
if not kv_connector_output:
- return None # noqa
+ return None # type: ignore[return-value]
# In case of PP with kv transfer, we need to pass through the
# kv_connector_output
@@ -2917,9 +2917,7 @@ class GPUModelRunner(
self.input_batch.prev_sampled_token_ids = None
- def propose_draft_token_ids(
- sampled_token_ids: torch.Tensor | list[np.ndarray],
- ) -> None:
+ def propose_draft_token_ids(sampled_token_ids):
assert spec_decode_common_attn_metadata is not None
with record_function_or_nullcontext("gpu_model_runner: draft"):
self._draft_token_ids = self.propose_draft_token_ids(
@@ -2933,33 +2931,37 @@ class GPUModelRunner(
spec_decode_common_attn_metadata,
)
+ spec_config = self.speculative_config
use_padded_batch_for_eagle = (
- self.speculative_config
- and self.speculative_config.use_eagle()
- and not self.speculative_config.disable_padded_drafter_batch
+ spec_config is not None
+ and spec_config.use_eagle()
+ and not spec_config.disable_padded_drafter_batch
)
effective_drafter_max_model_len = self.max_model_len
if effective_drafter_max_model_len is None:
effective_drafter_max_model_len = self.model_config.max_model_len
if (
- self.speculative_config
- and self.speculative_config.draft_model_config is not None
- and self.speculative_config.draft_model_config.max_model_len is not None
+ spec_config is not None
+ and spec_config.draft_model_config is not None
+ and spec_config.draft_model_config.max_model_len is not None
):
effective_drafter_max_model_len = (
- self.speculative_config.draft_model_config.max_model_len
+ spec_config.draft_model_config.max_model_len
)
input_fits_in_drafter = spec_decode_common_attn_metadata and (
spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
<= effective_drafter_max_model_len
)
if use_padded_batch_for_eagle:
+ assert self.speculative_config is not None
+ assert isinstance(self.drafter, EagleProposer)
sampled_token_ids = sampler_output.sampled_token_ids
if input_fits_in_drafter:
# EAGLE speculative decoding can use the GPU sampled tokens
# as inputs, and does not need to wait for bookkeeping to finish.
propose_draft_token_ids(sampled_token_ids)
elif self.valid_sampled_token_count_event is not None:
+ assert spec_decode_common_attn_metadata is not None
next_token_ids, valid_sampled_tokens_count = (
self.drafter.prepare_next_token_ids_padded(
spec_decode_common_attn_metadata,
@@ -3088,16 +3090,18 @@ class GPUModelRunner(
def propose_draft_token_ids(
self,
scheduler_output: "SchedulerOutput",
- sampled_token_ids: torch.Tensor | list[np.ndarray],
+ sampled_token_ids: torch.Tensor | list[list[int]],
sampling_metadata: SamplingMetadata,
hidden_states: torch.Tensor,
sample_hidden_states: torch.Tensor,
aux_hidden_states: list[torch.Tensor] | None,
spec_decode_metadata: SpecDecodeMetadata | None,
common_attn_metadata: CommonAttentionMetadata,
- ) -> torch.Tensor | list[list[int]]:
+ ) -> list[list[int]] | torch.Tensor:
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
- if self.speculative_config.method == "ngram":
+ spec_config = self.speculative_config
+ assert spec_config is not None
+ if spec_config.method == "ngram":
assert isinstance(sampled_token_ids, list)
assert isinstance(self.drafter, NgramProposer)
draft_token_ids = self.drafter.propose(
@@ -3107,11 +3111,11 @@ class GPUModelRunner(
self.input_batch.token_ids_cpu,
self.input_batch.spec_decode_unsupported_reqs,
)
- elif self.speculative_config.method == "suffix":
+ elif spec_config.method == "suffix":
assert isinstance(sampled_token_ids, list)
assert isinstance(self.drafter, SuffixDecodingProposer)
draft_token_ids = self.drafter.propose(self.input_batch, sampled_token_ids)
- elif self.speculative_config.method == "medusa":
+ elif spec_config.method == "medusa":
assert isinstance(sampled_token_ids, list)
assert isinstance(self.drafter, MedusaProposer)
@@ -3127,7 +3131,7 @@ class GPUModelRunner(
for num_draft, tokens in zip(
spec_decode_metadata.num_draft_tokens, sampled_token_ids
):
- indices.append(offset + tokens.shape[0] - 1)
+ indices.append(offset + len(tokens) - 1)
offset += num_draft + 1
indices = torch.tensor(indices, device=self.device)
hidden_states = sample_hidden_states[indices]
@@ -3136,10 +3140,10 @@ class GPUModelRunner(
target_hidden_states=hidden_states,
sampling_metadata=sampling_metadata,
)
- elif self.speculative_config.use_eagle():
+ elif spec_config.use_eagle():
assert isinstance(self.drafter, EagleProposer)
- if self.speculative_config.disable_padded_drafter_batch:
+ if spec_config.disable_padded_drafter_batch:
# When padded-batch is disabled, the sampled_token_ids should be
# the cpu-side list[list[int]] of valid sampled tokens for each
# request, with invalid requests having empty lists.
@@ -3189,7 +3193,7 @@ class GPUModelRunner(
else:
target_hidden_states = hidden_states[:num_scheduled_tokens]
else:
- if self.speculative_config.disable_padded_drafter_batch:
+ if spec_config.disable_padded_drafter_batch:
token_indices_to_sample = None
common_attn_metadata, token_indices = self.drafter.prepare_inputs(
common_attn_metadata,
@@ -3284,9 +3288,12 @@ class GPUModelRunner(
and is_mixture_of_experts(self.drafter.model)
and self.parallel_config.enable_eplb
):
+ spec_config = self.vllm_config.speculative_config
+ assert spec_config is not None
+ assert spec_config.draft_model_config is not None
logger.info_once(
"EPLB is enabled for drafter model %s.",
- self.vllm_config.speculative_config.draft_model_config.model,
+ spec_config.draft_model_config.model,
)
global_expert_load = (
@@ -3303,7 +3310,7 @@ class GPUModelRunner(
self.eplb_state = EplbState(self.parallel_config, self.device)
self.eplb_state.add_model(
self.drafter.model,
- self.vllm_config.speculative_config.draft_model_config,
+ spec_config.draft_model_config,
global_expert_load,
old_global_expert_indices,
rank_mapping,
@@ -3338,9 +3345,11 @@ class GPUModelRunner(
scope="local",
)
prepare_communication_buffer_for_model(self.model)
+ mm_config = self.model_config.multimodal_config
self.is_multimodal_pruning_enabled = (
supports_multimodal_pruning(self.get_model())
- and self.model_config.multimodal_config.is_multimodal_pruning_enabled()
+ and mm_config is not None
+ and mm_config.is_multimodal_pruning_enabled()
)
if is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb:
@@ -3375,15 +3384,14 @@ class GPUModelRunner(
# CudagraphWraper and CudagraphDispatcher of vllm.
# wrap the model with full cudagraph wrapper if needed.
- if (
- self.compilation_config.cudagraph_mode.has_full_cudagraphs()
- and not self.parallel_config.enable_dbo
- ):
+ cudagraph_mode = self.compilation_config.cudagraph_mode
+ assert cudagraph_mode is not None
+ if cudagraph_mode.has_full_cudagraphs() and not self.parallel_config.enable_dbo:
self.model = CUDAGraphWrapper(
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
)
elif self.parallel_config.enable_dbo:
- if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
+ if cudagraph_mode.has_full_cudagraphs():
self.model = UBatchWrapper(
self.model, self.vllm_config, CUDAGraphMode.FULL, self.device
)
@@ -3439,7 +3447,7 @@ class GPUModelRunner(
hidden_states: torch.Tensor,
num_scheduled_tokens: dict[str, int],
) -> dict[str, LogprobsTensors | None]:
- num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs
+ num_prompt_logprobs_dict = self.num_prompt_logprobs
if not num_prompt_logprobs_dict:
return {}
@@ -3450,7 +3458,10 @@ class GPUModelRunner(
# maintainable loop over optimal performance.
completed_prefill_reqs = []
for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items():
- num_tokens = num_scheduled_tokens[req_id]
+ num_tokens = num_scheduled_tokens.get(req_id)
+ if num_tokens is None:
+ # This can happen if the request was preempted in prefill stage.
+ continue
# Get metadata for this request.
request = self.requests[req_id]
@@ -3729,6 +3740,31 @@ class GPUModelRunner(
dp_rank = self.parallel_config.data_parallel_rank
num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])
+ # filter out the valid batch descriptor
+ _cg_mode, batch_descriptor = (
+ self.cudagraph_dispatcher.dispatch(
+ BatchDescriptor(
+ num_tokens=num_tokens_after_padding,
+ uniform_decode=uniform_decode,
+ has_lora=activate_lora and self.lora_config is not None,
+ )
+ )
+ if not is_profile
+ else (CUDAGraphMode.NONE, None)
+ )
+ if cudagraph_runtime_mode is not None:
+ # we allow forcing NONE when the dispatcher disagrees to support
+ # warm ups for cudagraph capture
+ assert (
+ cudagraph_runtime_mode == CUDAGraphMode.NONE
+ or cudagraph_runtime_mode == _cg_mode
+ ), (
+ f"Cudagraph runtime mode mismatch at dummy_run. "
+ f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+ )
+ else:
+ cudagraph_runtime_mode = _cg_mode
+
attn_metadata: PerLayerAttnMetadata | None = None
# If force_attention is True, we always capture attention. Otherwise,
@@ -3803,31 +3839,6 @@ class GPUModelRunner(
num_tokens_after_padding, None, False
)
- # filter out the valid batch descriptor
- _cg_mode, batch_descriptor = (
- self.cudagraph_dispatcher.dispatch(
- BatchDescriptor(
- num_tokens=num_tokens_after_padding,
- uniform_decode=uniform_decode,
- has_lora=activate_lora and self.lora_config is not None,
- )
- )
- if not is_profile
- else (CUDAGraphMode.NONE, None)
- )
- if cudagraph_runtime_mode is not None:
- # we allow forcing NONE when the dispatcher disagrees to support
- # warm ups for cudagraph capture
- assert (
- cudagraph_runtime_mode == CUDAGraphMode.NONE
- or cudagraph_runtime_mode == _cg_mode
- ), (
- f"Cudagraph runtime mode mismatch at dummy_run. "
- f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
- )
- else:
- cudagraph_runtime_mode = _cg_mode
-
if ubatch_slices is not None:
# Adjust values to reflect a single ubatch.
# TODO(sage,lucas): this is cruft that should be addressed in
@@ -4063,7 +4074,8 @@ class GPUModelRunner(
def profile_run(self) -> None:
# Profile with multimodal encoder & encoder cache.
if self.supports_mm_inputs:
- if self.model_config.multimodal_config.skip_mm_profiling:
+ mm_config = self.model_config.multimodal_config
+ if mm_config is not None and mm_config.skip_mm_profiling:
logger.info(
"Skipping memory profiling for multimodal encoder and "
"encoder cache."
@@ -4325,8 +4337,9 @@ class GPUModelRunner(
def get_attn_backends_for_group(
kv_cache_group_spec: KVCacheGroupSpec,
) -> tuple[dict[AttentionGroupKey, list[str]], set[type[AttentionBackend]]]:
+ layer_type = cast(type[Any], AttentionLayerBase)
layers = get_layers_from_vllm_config(
- self.vllm_config, AttentionLayerBase, kv_cache_group_spec.layer_names
+ self.vllm_config, layer_type, kv_cache_group_spec.layer_names
)
attn_backends = {}
attn_backend_layers = defaultdict(list)
@@ -4341,7 +4354,7 @@ class GPUModelRunner(
if layer_name in self.kv_sharing_fast_prefill_eligible_layers:
attn_backend = create_fast_prefill_custom_backend(
"FastPrefill",
- attn_backend,
+ attn_backend, # type: ignore[arg-type]
)
full_cls_name = attn_backend.full_cls_name()
@@ -4440,6 +4453,7 @@ class GPUModelRunner(
min_cg_backend_name = attn_backend.__name__
# Flexible resolve the cudagraph mode
cudagraph_mode = self.compilation_config.cudagraph_mode
+ assert cudagraph_mode is not None
# check cudagraph for mixed batch is supported
if (
cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL
@@ -4554,12 +4568,17 @@ class GPUModelRunner(
self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(
self.uniform_decode_query_len, self.parallel_config.tensor_parallel_size
)
- self.cudagraph_batch_sizes = self.compilation_config.cudagraph_capture_sizes
+ capture_sizes = self.compilation_config.cudagraph_capture_sizes
+ self.cudagraph_batch_sizes = (
+ capture_sizes if capture_sizes is not None else []
+ )
# Trigger cudagraph dispatching keys initialization after
# resolved cudagraph mode.
+ cudagraph_mode = self.compilation_config.cudagraph_mode
+ assert cudagraph_mode is not None
self.cudagraph_dispatcher.initialize_cudagraph_keys(
- self.compilation_config.cudagraph_mode, self.uniform_decode_query_len
+ cudagraph_mode, self.uniform_decode_query_len
)
def calculate_reorder_batch_threshold(self) -> None:
@@ -4571,7 +4590,7 @@ class GPUModelRunner(
"""
min_none_high = lambda a, b: a if b is None else b if a is None else min(a, b)
- reorder_batch_thresholds = [
+ reorder_batch_thresholds: list[int | None] = [
group.get_metadata_builder().reorder_batch_threshold
for group in self._attn_group_iterator()
]
@@ -4580,7 +4599,7 @@ class GPUModelRunner(
if len(reorder_batch_thresholds) == 0:
self.reorder_batch_threshold = None
return
- self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds)
+ self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment]
@staticmethod
def select_common_block_size(
@@ -4612,7 +4631,7 @@ class GPUModelRunner(
"""
for backend in backends:
is_supported = False
- for supported_size in backend.supported_kernel_block_sizes:
+ for supported_size in backend.get_supported_kernel_block_sizes():
if isinstance(supported_size, int):
if block_size == supported_size:
is_supported = True
@@ -4643,7 +4662,7 @@ class GPUModelRunner(
all_int_supported_sizes = set(
supported_size
for backend in backends
- for supported_size in backend.supported_kernel_block_sizes
+ for supported_size in backend.get_supported_kernel_block_sizes()
if isinstance(supported_size, int)
)
@@ -4925,12 +4944,30 @@ class GPUModelRunner(
Dict[str, torch.Tensor]: A map between layer names to their
corresponding memory buffer for KV cache.
"""
- # Initialize the memory buffer for KV cache
- kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
- # Change the memory buffer to the desired shape
- kv_caches = self._reshape_kv_cache_tensors(
- kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes
- )
+
+ # Try creating KV caches optimized for kv-connector transfers
+ cache_dtype = self.cache_config.cache_dtype
+ if self.use_uniform_kv_cache(self.attn_groups, cache_dtype):
+ kv_caches, cross_layers_kv_cache, attn_backend = (
+ self.allocate_uniform_kv_caches(
+ kv_cache_config,
+ self.attn_groups,
+ cache_dtype,
+ self.device,
+ kernel_block_sizes,
+ )
+ )
+ self.cross_layers_kv_cache = cross_layers_kv_cache
+ self.cross_layers_attn_backend = attn_backend
+ else:
+ # Fallback to the general case
+ # Initialize the memory buffer for KV cache
+ kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
+
+ # Change the memory buffer to the desired shape
+ kv_caches = self._reshape_kv_cache_tensors(
+ kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes
+ )
# Set up cross-layer KV cache sharing
for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
@@ -5012,16 +5049,26 @@ class GPUModelRunner(
if has_kv_transfer_group():
kv_transfer_group = get_kv_transfer_group()
- kv_transfer_group.register_kv_caches(kv_caches)
+ if self.cross_layers_kv_cache is not None:
+ assert self.cross_layers_attn_backend is not None
+ kv_transfer_group.register_cross_layers_kv_cache(
+ self.cross_layers_kv_cache, self.cross_layers_attn_backend
+ )
+ else:
+ kv_transfer_group.register_kv_caches(kv_caches)
kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
if self.dcp_world_size > 1:
- layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
+ layer_type = cast(type[Any], AttentionLayerBase)
+ layers = get_layers_from_vllm_config(self.vllm_config, layer_type)
for layer in layers.values():
- assert layer.impl.need_to_return_lse_for_decode, (
+ layer_impl = getattr(layer, "impl", None)
+ if layer_impl is None:
+ continue
+ assert layer_impl.need_to_return_lse_for_decode, (
"DCP requires attention impls to return"
" the softmax lse for decode, but the impl "
- f"{layer.impl.__class__.__name__} "
+ f"{layer_impl.__class__.__name__} "
"does not return the softmax lse for decode."
)
@@ -5062,7 +5109,8 @@ class GPUModelRunner(
if has_ec_transfer() and get_ec_transfer().is_producer:
return {}
kv_cache_spec: dict[str, KVCacheSpec] = {}
- attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
+ layer_type = cast(type[Any], AttentionLayerBase)
+ attn_layers = get_layers_from_vllm_config(self.vllm_config, layer_type)
for layer_name, attn_module in attn_layers.items():
if isinstance(attn_module, Attention) and (
kv_tgt_layer := attn_module.kv_sharing_target_layer_name
@@ -5082,7 +5130,7 @@ class GPUModelRunner(
return kv_cache_spec
- def _to_list(self, sampled_token_ids: torch.Tensor) -> list[np.ndarray]:
+ def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
# This is a short term mitigation for issue mentioned in
# https://github.com/vllm-project/vllm/issues/22754.
# `tolist` would trigger a cuda wise stream sync, which
@@ -5095,4 +5143,4 @@ class GPUModelRunner(
pinned.copy_(sampled_token_ids, non_blocking=True)
self.transfer_event.record()
self.transfer_event.synchronize()
- return [row for row in pinned.numpy()]
+ return pinned.tolist()
diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py
index 9de123263755b..2ce2b64512560 100644
--- a/vllm/v1/worker/gpu_ubatch_wrapper.py
+++ b/vllm/v1/worker/gpu_ubatch_wrapper.py
@@ -121,18 +121,24 @@ class UBatchWrapper:
@staticmethod
def _create_sm_control_context(vllm_config: VllmConfig):
- comm_sms = envs.VLLM_DBO_COMM_SMS
+ comm_sms: int = envs.VLLM_DBO_COMM_SMS
set_comm_sms = lambda sms: None
if vllm_config.parallel_config.enable_expert_parallel:
# Currently only DeepEP highthroughput supports SM control so this
# only affects that case.
- all2all_manager = get_ep_group().device_communicator.all2all_manager
+ ep_group = get_ep_group()
+ device_communicator = ep_group.device_communicator
+ all2all_manager = None
+ if device_communicator is not None:
+ all2all_manager = device_communicator.all2all_manager
- if all2all_manager.max_sms_used() is not None:
- comm_sms = min(comm_sms, all2all_manager.max_sms_used())
+ if all2all_manager is not None:
+ max_sms_used = all2all_manager.max_sms_used()
+ if max_sms_used is not None:
+ comm_sms = min(comm_sms, max_sms_used)
- if comm_sms > 0:
+ if comm_sms > 0 and all2all_manager is not None:
set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms)
# TODO(lucas): support other kernels besides DeepGEMM
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index 315f01b68499a..6a4bfde5f972b 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -6,7 +6,7 @@ import gc
import os
from contextlib import AbstractContextManager, nullcontext
from types import NoneType
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
import torch
import torch.distributed
@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer import (
has_kv_transfer_group,
)
from vllm.distributed.parallel_state import (
+ get_pcp_group,
get_pp_group,
get_tp_group,
)
@@ -35,12 +36,12 @@ from vllm.model_executor import set_random_seed
from vllm.model_executor.models.interfaces import is_mixture_of_experts
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
from vllm.platforms import current_platform
-from vllm.profiler.gpu_profiler import CudaProfilerWrapper
+from vllm.profiler.gpu_profiler import CudaProfilerWrapper, TorchProfilerWrapper
from vllm.sequence import IntermediateTensors
from vllm.tasks import SupportedTask
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
-from vllm.v1.core.sched.output import GrammarOutput
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm.v1.outputs import (
@@ -57,7 +58,6 @@ logger = init_logger(__name__)
if TYPE_CHECKING:
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
- from vllm.v1.core.sched.output import SchedulerOutput
class Worker(WorkerBase):
@@ -86,41 +86,22 @@ class Worker(WorkerBase):
# Buffers saved before sleep
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
- # Torch profiler. Enabled and configured through env vars:
+ # Torch/CUDA profiler. Enabled and configured through env vars:
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+ # VLLM_TORCH_CUDA_PROFILE=1
+ self.profiler: Any | None = None
if envs.VLLM_TORCH_PROFILER_DIR:
- torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
- logger.info(
- "Profiling enabled. Traces will be saved to: %s",
- torch_profiler_trace_dir,
- )
- logger.debug(
- "Profiler config: record_shapes=%s,"
- "profile_memory=%s,with_stack=%s,with_flops=%s",
- envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
- envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
- envs.VLLM_TORCH_PROFILER_WITH_STACK,
- envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
- )
- self.profiler = torch.profiler.profile(
- activities=[
- torch.profiler.ProfilerActivity.CPU,
- torch.profiler.ProfilerActivity.CUDA,
- ],
- record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
- profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
- with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
- with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
- torch_profiler_trace_dir, worker_name=worker_name, use_gzip=True
- ),
+ self.profiler = TorchProfilerWrapper(
+ worker_name=worker_name, local_rank=self.local_rank
)
elif envs.VLLM_TORCH_CUDA_PROFILE:
self.profiler = CudaProfilerWrapper()
else:
self.profiler = None
+ self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
+
def sleep(self, level: int = 1) -> None:
from vllm.device_allocator.cumem import CuMemAllocator
@@ -168,17 +149,17 @@ class Worker(WorkerBase):
assert allocator.get_current_usage() == 0, (
"Sleep mode can only be used for one instance per process."
)
- context = allocator.use_memory_pool(tag=tag)
+ return allocator.use_memory_pool(tag=tag)
else:
- context = nullcontext()
- return context
+ return nullcontext()
def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None:
self.cache_config.num_gpu_blocks = num_gpu_blocks
self.cache_config.num_cpu_blocks = num_cpu_blocks
def init_device(self):
- if self.device_config.device.type == "cuda":
+ device = self.device_config.device
+ if isinstance(device, torch.device) and device.type == "cuda":
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
if (
@@ -204,14 +185,14 @@ class Worker(WorkerBase):
assert self.local_rank < torch.cuda.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
- visible_device_count = (
- torch.cuda.device_count() if torch.cuda.is_available() else 0
- )
- assert self.parallel_config.local_world_size <= visible_device_count, (
- f"local_world_size ({self.parallel_config.local_world_size}) must be "
- f"less than or equal to the number of visible devices "
- f"({visible_device_count})."
- )
+ visible_device_count = (
+ torch.cuda.device_count() if torch.cuda.is_available() else 0
+ )
+ assert self.parallel_config.local_world_size <= visible_device_count, (
+ f"local_world_size ({self.parallel_config.local_world_size}) must "
+ f"be less than or equal to the number of visible devices "
+ f"({visible_device_count})."
+ )
self.device = torch.device(f"cuda:{self.local_rank}")
current_platform.set_device(self.device)
@@ -257,9 +238,17 @@ class Worker(WorkerBase):
raise RuntimeError(f"Not support device type: {self.device_config.device}")
# Construct the model runner
- self.model_runner: GPUModelRunner = GPUModelRunner(
- self.vllm_config, self.device
- )
+ if self.use_v2_model_runner:
+ from vllm.v1.worker.gpu.model_runner import (
+ GPUModelRunner as GPUModelRunnerV2,
+ )
+
+ # HACK(woosuk): This is a temporary fix to avoid type errors.
+ self.model_runner: GPUModelRunner = GPUModelRunnerV2( # type: ignore
+ self.vllm_config, self.device
+ )
+ else:
+ self.model_runner = GPUModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
@@ -397,23 +386,21 @@ class Worker(WorkerBase):
from vllm.device_allocator.cumem import CuMemAllocator
allocator = CuMemAllocator.get_instance()
- context = allocator.use_memory_pool(tag="kv_cache")
+ with allocator.use_memory_pool(tag="kv_cache"):
+ self.model_runner.initialize_kv_cache(kv_cache_config)
else:
- context = nullcontext()
- with context:
self.model_runner.initialize_kv_cache(kv_cache_config)
def compile_or_warm_up_model(self) -> None:
# warm up sizes that are not in cudagraph capture sizes,
# but users still want to compile for better performance,
# e.g. for the max-num-batched token size in chunked prefill.
- warmup_sizes = self.vllm_config.compilation_config.compile_sizes.copy()
+ compile_sizes = self.vllm_config.compilation_config.compile_sizes
+ warmup_sizes = compile_sizes.copy() if compile_sizes is not None else []
if not self.model_config.enforce_eager:
- warmup_sizes = [
- x
- for x in warmup_sizes
- if x not in self.vllm_config.compilation_config.cudagraph_capture_sizes
- ]
+ capture_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
+ if capture_sizes is not None:
+ warmup_sizes = [x for x in warmup_sizes if x not in capture_sizes]
# We skip EPLB here since we don't want to record dummy metrics
for size in sorted(warmup_sizes, reverse=True):
logger.info("Compile and warming up model for size %d", size)
@@ -525,10 +512,12 @@ class Worker(WorkerBase):
if not self.profiler:
return nullcontext()
+ self.profiler.step()
+
num_new = len(scheduler_output.scheduled_new_reqs)
num_cached = len(scheduler_output.scheduled_cached_reqs.req_ids)
- return torch.profiler.record_function(
+ return self.profiler.annotate_context_manager(
f"execute_new_{num_new}_cached_{num_cached}"
)
@@ -552,12 +541,12 @@ class Worker(WorkerBase):
)
}
if forward_pass and not get_pp_group().is_first_rank:
- intermediate_tensors = IntermediateTensors(
- get_pp_group().recv_tensor_dict(
- all_gather_group=get_tp_group(),
- all_gather_tensors=all_gather_tensors,
- )
+ tensor_dict = get_pp_group().recv_tensor_dict(
+ all_gather_group=get_tp_group(),
+ all_gather_tensors=all_gather_tensors,
)
+ assert tensor_dict is not None
+ intermediate_tensors = IntermediateTensors(tensor_dict)
with self.annotate_profile(scheduler_output):
output = self.model_runner.execute_model(
@@ -586,27 +575,19 @@ class Worker(WorkerBase):
def profile(self, is_start: bool = True):
if self.profiler is None:
- raise RuntimeError("Profiler is not enabled.")
+ raise RuntimeError("Profiling is not enabled.")
if is_start:
self.profiler.start()
else:
self.profiler.stop()
- if isinstance(self.profiler, torch.profiler.profile):
- rank = self.local_rank
- profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
- profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
- sort_key = "self_cuda_time_total"
- table = self.profiler.key_averages().table(sort_by=sort_key)
-
- with open(profiler_out_file, "w") as f:
- print(table, file=f)
-
- # only print profiler results on rank 0
- if rank == 0:
- print(table)
def execute_dummy_batch(self) -> None:
- self.model_runner._dummy_run(1, uniform_decode=True)
+ if self.use_v2_model_runner:
+ self.model_runner.execute_model(
+ SchedulerOutput.make_empty(), dummy_run=True
+ )
+ else:
+ self.model_runner._dummy_run(1, uniform_decode=True)
def add_lora(self, lora_request: LoRARequest) -> bool:
return self.model_runner.add_lora(lora_request)
@@ -638,7 +619,7 @@ class Worker(WorkerBase):
assert self.model_runner.eplb_state is not None
self.model_runner.eplb_state.rearrange(
execute_shuffle=True,
- global_expert_load=None,
+ global_expert_loads=None,
rank_mapping=rank_mapping,
)
torch.cuda.synchronize()
@@ -694,7 +675,7 @@ class Worker(WorkerBase):
def _reconfigure_moe(
self, old_ep_size: int, new_ep_size: int
- ) -> torch.Tensor | None:
+ ) -> list[torch.Tensor] | None:
"""
Reconfigure MoE modules with provided reconfig_request
@@ -733,6 +714,7 @@ class Worker(WorkerBase):
module.global_num_experts = module.moe_config.num_experts
module.moe_parallel_config = FusedMoEParallelConfig.make(
tp_size_=get_tp_group().world_size,
+ pcp_size_=get_pcp_group().world_size,
dp_size_=get_dp_group().world_size,
vllm_parallel_config=parallel_config,
)
@@ -760,26 +742,29 @@ class Worker(WorkerBase):
num_local_physical_experts = num_local_experts
assert self.model_runner.eplb_state is not None
new_physical_experts = (
- self.model_runner.eplb_state.physical_to_logical_map.shape[1]
+ self.model_runner.eplb_state.physical_to_logical_map.shape[1] # type: ignore[attr-defined]
)
parallel_config.eplb_config.num_redundant_experts = (
new_physical_experts
- - self.model_runner.eplb_state.logical_replica_count.shape[1]
+ - self.model_runner.eplb_state.logical_replica_count.shape[1] # type: ignore[attr-defined]
)
global_expert_loads = None
else:
- num_local_physical_experts = torch.tensor(
+ num_local_physical_experts_tensor = torch.tensor(
[num_local_experts], dtype=torch.int32, device="cpu"
)
torch.distributed.broadcast(
- num_local_physical_experts, group=get_ep_group().cpu_group, group_src=0
+ num_local_physical_experts_tensor,
+ group=get_ep_group().cpu_group,
+ group_src=0,
)
- num_local_physical_experts = num_local_physical_experts.item()
+ num_local_physical_experts = int(num_local_physical_experts_tensor.item())
new_physical_experts = num_local_physical_experts * new_ep_size
assert self.model_runner.eplb_state is not None
- global_expert_loads = self.model_runner.eplb_state.rearrange(
+ global_expert_loads_any = self.model_runner.eplb_state.rearrange(
execute_shuffle=False
)
+ global_expert_loads = cast(list[torch.Tensor], global_expert_loads_any)
parallel_config.eplb_config.num_redundant_experts = (
new_physical_experts - global_expert_loads[0].shape[1]
)
@@ -863,6 +848,8 @@ class Worker(WorkerBase):
def shutdown(self) -> None:
if runner := getattr(self, "model_runner", None):
runner.ensure_kv_transfer_shutdown()
+ if self.profiler is not None:
+ self.profiler.shutdown()
def init_worker_distributed_environment(
@@ -879,13 +866,15 @@ def init_worker_distributed_environment(
init_batch_invariance()
set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
+ init_method = distributed_init_method or "env://"
init_distributed_environment(
- parallel_config.world_size, rank, distributed_init_method, local_rank, backend
+ parallel_config.world_size, rank, init_method, local_rank, backend
)
ensure_model_parallel_initialized(
parallel_config.tensor_parallel_size,
parallel_config.pipeline_parallel_size,
+ parallel_config.prefill_context_parallel_size,
parallel_config.decode_context_parallel_size,
)
diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py
index db037a9fccd5c..ff047d8d03f0e 100644
--- a/vllm/v1/worker/kv_connector_model_runner_mixin.py
+++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py
@@ -11,7 +11,11 @@ from typing import (
TYPE_CHECKING, # noqa: UP035
)
+import torch
+
+from vllm.attention import AttentionBackend
from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
from vllm.distributed.kv_transfer import (
ensure_kv_transfer_shutdown,
get_kv_transfer_group,
@@ -21,11 +25,13 @@ from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
from vllm.forward_context import get_forward_context, set_forward_context
from vllm.logger import init_logger
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
from vllm.v1.outputs import (
EMPTY_MODEL_RUNNER_OUTPUT,
KVConnectorOutput,
ModelRunnerOutput,
)
+from vllm.v1.worker.utils import AttentionGroup
if TYPE_CHECKING:
from vllm.v1.core.sched.output import SchedulerOutput
@@ -53,7 +59,7 @@ class KVConnectorModelRunnerMixin:
@staticmethod
def ensure_kv_transfer_shutdown() -> None:
# has_kv_transfer_group can be None during interpreter shutdown.
- if has_kv_transfer_group and has_kv_transfer_group():
+ if has_kv_transfer_group and has_kv_transfer_group(): # type: ignore[truthy-function]
ensure_kv_transfer_shutdown()
@staticmethod
@@ -142,3 +148,162 @@ class KVConnectorModelRunnerMixin:
if has_kv_transfer_group():
return get_kv_transfer_group().get_kv_connector_stats()
return None
+
+ @staticmethod
+ def use_uniform_kv_cache(
+ attn_groups: list[list[AttentionGroup]],
+ cache_dtype: CacheDType,
+ ) -> bool:
+ """
+ Determines whether a uniform KV layout should be used.
+ A uniform layout means all layers KV caches will share the same
+ underlying tensor, where for a given block number, the respective
+ KV data for all layers will be contiguous.
+ This will allow efficient KV transfer of per-block KV data for all
+ layers at once.
+ Note this layout will only be applied given 3 conditions:
+ 1. The KV Cache config contains just a single group where all layers
+ have the same page size.
+ 2. A KV connector is configured, and the KV connector instance prefers
+ to use this layout (prefer_cross_layer_blocks() returns True)
+ 2. The flash attention backend supports this layout
+ (get_kv_cache_stride_order(True) includes a placement for a
+ num_layers dimension)
+
+ Note that the actual placement of the num_layers dimensions
+ in the unified layers tensors will be determined by the attention
+ backend.
+ Thus, the layers KV data may still not be contiguous per block
+ if the attention backend does not support it.
+
+ Args:
+ attn_groups: The list of attention groups for this model
+ cache_dtype: The KV cache dtype
+ Returns:
+ True if we should use a uniform KV cache layout.
+ """
+
+ if not has_kv_transfer_group():
+ return False
+ if not get_kv_transfer_group().prefer_cross_layer_blocks:
+ return False
+
+ if len(attn_groups) != 1 or len(attn_groups[0]) != 1:
+ return False
+
+ attn_group = attn_groups[0][0]
+ kv_cache_spec = attn_group.kv_cache_spec
+ if not isinstance(kv_cache_spec, AttentionSpec):
+ return False
+
+ attn_backend = attn_group.backend
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
+ 1234,
+ kv_cache_spec.block_size,
+ kv_cache_spec.num_kv_heads,
+ kv_cache_spec.head_size,
+ cache_dtype_str=cache_dtype,
+ )
+
+ try:
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
+ include_num_layers_dimension=True
+ )
+ except (AttributeError, NotImplementedError):
+ return False
+
+ # check that attention backend include a layers dimension
+ return len(kv_cache_stride_order) == len(kv_cache_shape) + 1
+
+ @staticmethod
+ def allocate_uniform_kv_caches(
+ kv_cache_config: KVCacheConfig,
+ attn_groups: list[list[AttentionGroup]],
+ cache_dtype: CacheDType,
+ device: torch.device,
+ kernel_block_sizes: list[int],
+ ) -> tuple[dict[str, torch.Tensor], torch.Tensor, type[AttentionBackend]]:
+ """
+ Initializes and reshapes KV caches for the simple case where all
+ layers have the same layout.
+
+ This function assumes use_uniform_kv_cache() returned True.
+
+ Args:
+ kv_cache_config: The KV cache config
+ attn_groups: The list of attention groups for this model
+ cache_dtype: The KV cache dtype
+ device: The torch device to allocate on.
+ kernel_block_sizes: The kernel block sizes for each KV cache group.
+ Returns:
+ A tuple (kv_caches, cross_layers_kv_cache, attn_backend) where:
+ kv_caches is a dict mapping between layer names to their
+ corresponding memory buffer for KV cache.
+ cross_layers_kv_cache is the cross layers kv cache tensor
+ attn_backend is the attention backend matching this tensor
+ """
+ attn_group = attn_groups[0][0]
+ kv_cache_spec = attn_group.kv_cache_spec
+ assert isinstance(kv_cache_spec, AttentionSpec)
+
+ tensor_sizes = set(
+ kv_cache_tensor.size for kv_cache_tensor in kv_cache_config.kv_cache_tensors
+ )
+ assert len(tensor_sizes) == 1
+ tensor_size = tensor_sizes.pop()
+
+ page_size = kv_cache_spec.page_size_bytes
+ assert tensor_size % page_size == 0
+ num_blocks = tensor_size // page_size
+ num_layers = len(kv_cache_config.kv_cache_tensors)
+ total_size = tensor_size * num_layers
+
+ assert len(kernel_block_sizes) == 1
+ kernel_block_size = kernel_block_sizes[0]
+ num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size
+ kernel_num_blocks = num_blocks * num_blocks_per_kv_block
+
+ attn_backend = attn_group.backend
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
+ kernel_num_blocks,
+ kernel_block_size,
+ kv_cache_spec.num_kv_heads,
+ kv_cache_spec.head_size,
+ cache_dtype_str=cache_dtype,
+ )
+
+ # prepend a num_layers dimension into the shape
+ kv_cache_shape = (num_layers,) + kv_cache_shape
+
+ try:
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
+ include_num_layers_dimension=True
+ )
+ assert len(kv_cache_stride_order) == len(kv_cache_shape)
+ except (AttributeError, NotImplementedError):
+ kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+ kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+
+ logger.info("Allocating a cross layer KV cache of shape %s", kv_cache_shape)
+
+ # allocate one contiguous buffer for all layers
+ cross_layers_kv_cache = (
+ torch.zeros(total_size, dtype=torch.int8, device=device)
+ .view(kv_cache_spec.dtype)
+ .view(kv_cache_shape)
+ )
+
+ # Maintain original KV shape view.
+ inv_order = [
+ kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+ ]
+ permuted_kv_cache = cross_layers_kv_cache.permute(*inv_order)
+
+ kv_caches = {}
+ for i, kv_cache_tensor in enumerate(kv_cache_config.kv_cache_tensors):
+ tensor = permuted_kv_cache[i]
+ for layer_name in kv_cache_tensor.shared_by:
+ kv_caches[layer_name] = tensor
+
+ return kv_caches, cross_layers_kv_cache, attn_backend
diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py
index 6bf4f91931849..2ed65ca9d31cd 100644
--- a/vllm/v1/worker/tpu_input_batch.py
+++ b/vllm/v1/worker/tpu_input_batch.py
@@ -149,9 +149,6 @@ class InputBatch:
self.generators: dict[int, torch.Generator] = {}
self.num_logprobs: dict[str, int] = {}
- # NOTE(rob): num_prompt_logprobs only includes reqs
- # that are currently in the prefill phase.
- self.num_prompt_logprobs: dict[str, int] = {}
# To accumulate prompt logprobs tensor chunks across prefill steps.
self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
@@ -256,8 +253,6 @@ class InputBatch:
if sampling_params.logprobs is not None:
self.num_logprobs[req_id] = sampling_params.logprobs
- if sampling_params.prompt_logprobs is not None:
- self.num_prompt_logprobs[req_id] = sampling_params.prompt_logprobs
if sampling_params.logit_bias is not None:
self.logit_bias[req_index] = sampling_params.logit_bias
@@ -317,7 +312,6 @@ class InputBatch:
self.repetition_penalties_reqs.discard(req_id)
self.generators.pop(req_index, None)
self.num_logprobs.pop(req_id, None)
- self.num_prompt_logprobs.pop(req_id, None)
self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
# LoRA
@@ -584,10 +578,6 @@ class InputBatch:
def max_num_logprobs(self) -> int | None:
return max(self.num_logprobs.values()) if self.num_logprobs else None
- @property
- def no_prompt_logprob(self) -> bool:
- return not self.num_prompt_logprobs
-
@property
def no_allowed_token_ids(self) -> bool:
return len(self.has_allowed_token_ids) == 0
diff --git a/vllm/v1/worker/tpu_model_runner.py b/vllm/v1/worker/tpu_model_runner.py
index e9eb7cad38f88..72d4474b89627 100644
--- a/vllm/v1/worker/tpu_model_runner.py
+++ b/vllm/v1/worker/tpu_model_runner.py
@@ -219,9 +219,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
self.hidden_size = model_config.get_hidden_size()
self.vocab_size = model_config.get_vocab_size()
- if self.lora_config is not None:
- self.vocab_size += self.lora_config.lora_extra_vocab_size
-
# Multi-modal data support
self.mm_registry = MULTIMODAL_REGISTRY
self.uses_mrope = model_config.uses_mrope
@@ -250,6 +247,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
# Request states.
self.requests: dict[str, CachedRequestState] = {}
+ # NOTE(rob): num_prompt_logprobs only includes reqs
+ # that are currently in the prefill phase.
+ self.num_prompt_logprobs: dict[str, int] = {}
# Initialize input batch early to avoid AttributeError in _update_states
self.input_batch = InputBatch(
@@ -423,6 +423,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
# Remove finished requests from the cached states.
for req_id in scheduler_output.finished_req_ids:
self.requests.pop(req_id, None)
+ self.num_prompt_logprobs.pop(req_id, None)
# Remove the finished requests from the persistent batch.
# NOTE(woosuk): There could be an edge case where finished_req_ids and
@@ -480,6 +481,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
lora_request=new_req_data.lora_request,
)
+ if sampling_params and sampling_params.prompt_logprobs is not None:
+ self.num_prompt_logprobs[req_id] = (
+ self.input_batch.vocab_size
+ if sampling_params.prompt_logprobs == -1
+ else sampling_params.prompt_logprobs
+ )
+
req_ids_to_add.append(req_id)
# Update the states of the running/resumed requests.
@@ -575,7 +583,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
format. Layers that do not need KV cache are not included.
"""
- layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
+ layers = get_layers_from_vllm_config(
+ self.vllm_config,
+ AttentionLayerBase, # type: ignore[type-abstract]
+ )
block_size = self.vllm_config.cache_config.block_size
cache_dtype_str = self.vllm_config.cache_config.cache_dtype
@@ -728,7 +739,11 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
req_id = self.input_batch.req_ids[i]
assert req_id is not None
num_tokens = scheduler_output.num_scheduled_tokens[req_id]
- if not use_max_model_len and num_tokens > self.most_model_len:
+ if (
+ not use_max_model_len
+ and self.most_model_len is not None
+ and num_tokens > self.most_model_len
+ ):
use_max_model_len = True
num_scheduled_tokens_per_req.append(num_tokens)
if use_max_model_len:
@@ -740,6 +755,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
else:
end_index = num_reqs
else:
+ assert self.num_reqs_most_model_len is not None
if len(num_scheduled_tokens_per_req) > self.num_reqs_most_model_len:
num_scheduled_tokens_per_req = num_scheduled_tokens_per_req[
: self.num_reqs_most_model_len
@@ -832,6 +848,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
].to(self.device)
seq_lens = self.seq_lens_cpu[: self.num_reqs_max_model_len].to(self.device)
else:
+ assert self.num_reqs_most_model_len is not None
block_tables = self.block_table_cpu[
: self.num_reqs_most_model_len, : self.num_blocks_per_most_len_req
]
@@ -934,6 +951,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
for mm_input_id in encoder_input_ids:
mm_feature = req_state.mm_features[mm_input_id]
+ if mm_feature.data is None:
+ continue
mm_hash = mm_feature.identifier
mm_kwargs.append(mm_feature.data)
mm_hashes_pos.append((mm_hash, mm_feature.mm_position))
@@ -1117,7 +1136,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
) -> ModelRunnerOutput:
if self.scheduler_output is None:
# Nothing to do (PP non-final rank case), output isn't used.
- return None # noqa
+ return None # type: ignore[return-value]
scheduler_output = self.scheduler_output
mm_embed_inputs = self.mm_embed_inputs
self.scheduler_output = None
@@ -1254,15 +1273,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
max_gen_len = selected_token_ids.shape[-1]
if max_gen_len == 1:
- valid_sampled_token_ids: list[np.ndarray] = [
- row for row in selected_token_ids.numpy()
- ]
+ valid_sampled_token_ids = selected_token_ids.tolist()
# Mask out the sampled tokens that should not be sampled.
# TODO: Keep in sync with gpu_model_runner.py, in particular
# the "else" case here
for i in discard_sampled_tokens_req_indices:
- valid_sampled_token_ids[i] = np.array([])
+ valid_sampled_token_ids[i].clear()
# Append sampled tokens
for i, req_state, seq_len in request_seq_lens:
@@ -1275,7 +1292,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
valid_mask = selected_token_ids != INVALID_TOKEN_ID
gen_lens = valid_mask.sum(dim=1).tolist()
valid_sampled_token_ids = [
- seq.numpy() for seq in selected_token_ids[valid_mask].split(gen_lens)
+ seq.tolist() for seq in selected_token_ids[valid_mask].split(gen_lens)
]
self.input_batch.num_tokens[:num_reqs] += gen_lens
for i, req_state, seq_len in request_seq_lens:
@@ -1699,7 +1716,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
) -> None:
# Profile with multimodal encoder & encoder cache.
if self.supports_mm_inputs:
- if self.model_config.multimodal_config.skip_mm_profiling:
+ mm_config = self.model_config.multimodal_config
+ if mm_config is not None and mm_config.skip_mm_profiling:
logger.info(
"Skipping memory profiling for multimodal encoder and "
"encoder cache."
@@ -2169,5 +2187,9 @@ def replace_set_lora(model):
if isinstance(module, BaseLayerWithLoRA):
module._original_set_lora = module.set_lora
module._original_reset_lora = module.reset_lora
- module.set_lora = _tpu_set_lora.__get__(module, module.__class__)
- module.reset_lora = _tpu_reset_lora.__get__(module, module.__class__)
+ module.set_lora = _tpu_set_lora.__get__( # type: ignore[method-assign]
+ module, module.__class__
+ )
+ module.reset_lora = _tpu_reset_lora.__get__( # type: ignore[method-assign]
+ module, module.__class__
+ )
diff --git a/vllm/v1/worker/tpu_worker.py b/vllm/v1/worker/tpu_worker.py
index a716a9c3aa822..e1a109eca0a88 100644
--- a/vllm/v1/worker/tpu_worker.py
+++ b/vllm/v1/worker/tpu_worker.py
@@ -106,9 +106,6 @@ class TPUWorker:
"Profiling enabled. Traces will be saved to: %s", self.profile_dir
)
- if self.model_config.seed is None:
- self.model_config.seed = 0
-
def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None:
self.cache_config.num_gpu_blocks = num_gpu_blocks
self.cache_config.num_cpu_blocks = num_cpu_blocks
@@ -141,8 +138,7 @@ class TPUWorker:
# Set random seed.
set_random_seed(self.model_config.seed)
- if self.model_config.seed is not None:
- xm.set_rng_state(self.model_config.seed, self.device)
+ xm.set_rng_state(self.model_config.seed, self.device)
# Increase the cache size limit, which is the maximum number of
# dynamo graphs that can be compiled.
@@ -332,7 +328,7 @@ class TPUWorker:
world_size=parallel_config.world_size,
rank=rank,
local_rank=local_rank,
- distributed_init_method=distributed_init_method,
+ distributed_init_method=distributed_init_method or "env://",
backend=current_platform.dist_backend,
)
ensure_model_parallel_initialized(
diff --git a/vllm/v1/worker/ubatching.py b/vllm/v1/worker/ubatching.py
index 9f16b1e6d03ee..be8326e2fdbc1 100644
--- a/vllm/v1/worker/ubatching.py
+++ b/vllm/v1/worker/ubatching.py
@@ -27,8 +27,8 @@ class UBatchContext:
ready_barrier: threading.Barrier,
cpu_wait_event: threading.Event,
cpu_signal_event: threading.Event,
- gpu_comm_done_event: torch.cuda.Event,
- gpu_compute_done_event: torch.cuda.Event,
+ gpu_comm_done_event: torch.Event,
+ gpu_compute_done_event: torch.Event,
schedule: str = "default",
):
self.id = id
@@ -207,8 +207,8 @@ def make_ubatch_contexts(
Create a context manager for micro-batching synchronization.
"""
cpu_events = [threading.Event() for _ in range(num_micro_batches)]
- gpu_comm_done_events = [torch.cuda.Event() for _ in range(num_micro_batches)]
- gpu_compute_done_events = [torch.cuda.Event() for _ in range(num_micro_batches)]
+ gpu_comm_done_events = [torch.Event() for _ in range(num_micro_batches)]
+ gpu_compute_done_events = [torch.Event() for _ in range(num_micro_batches)]
assert len(forward_contexts) == 2
diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py
index 095407a8b9596..92e4ce3abdba3 100644
--- a/vllm/v1/worker/utils.py
+++ b/vllm/v1/worker/utils.py
@@ -280,7 +280,7 @@ def bind_kv_cache(
kv_caches: dict[str, torch.Tensor],
forward_context: dict[str, "Attention"],
runner_kv_caches: list[torch.Tensor],
- num_attn_module: int | None = 1,
+ num_attn_module: int = 1,
) -> None:
"""
Bind the allocated KV cache to both ModelRunner and forward context so
@@ -316,7 +316,7 @@ def bind_kv_cache(
# TODO - analyze where runner_kv_caches is used and the right
# way to ensure it properly reflects multiple attention layers
# in the same decoder block.
- if current_platform.is_cuda() or current_platform.is_xpu():
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
# We know that the GPU runner is not impacted by this
# case. Some test code depends on runner_kv_caches, but
# not in a way that's impacted by ignoring this.
@@ -362,5 +362,7 @@ def is_residual_scattered_for_sp(
or vllm_config.compilation_config.use_inductor_graph_partition
):
return True
-
- return num_input_tokens in vllm_config.compilation_config.compile_sizes
+ compile_sizes = vllm_config.compilation_config.compile_sizes
+ if compile_sizes is None:
+ return False
+ return num_input_tokens in compile_sizes
diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py
index 16f321c080779..57e7037e946ec 100644
--- a/vllm/v1/worker/worker_base.py
+++ b/vllm/v1/worker/worker_base.py
@@ -315,10 +315,12 @@ class WorkerWrapperBase:
def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
kv_cache_config = kv_cache_configs[self.global_rank]
+ assert self.vllm_config is not None
with set_current_vllm_config(self.vllm_config):
self.worker.initialize_from_config(kv_cache_config) # type: ignore
def init_device(self):
+ assert self.vllm_config is not None
with set_current_vllm_config(self.vllm_config):
# To make vLLM config available during device initialization
self.worker.init_device() # type: ignore
diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py
index 4f82c18da73aa..30563305853a5 100644
--- a/vllm/v1/worker/xpu_model_runner.py
+++ b/vllm/v1/worker/xpu_model_runner.py
@@ -37,19 +37,12 @@ class XPUModelRunner(GPUModelRunner):
@contextmanager
def _torch_cuda_wrapper():
- class _EventPlaceholder:
- def __init__(self, *args, **kwargs) -> None:
- self.record = lambda: None
- self.synchronize = lambda: None
-
try:
# replace cuda APIs with xpu APIs, this should work by default
- torch.cuda.Event = torch.xpu.Event
torch.cuda.Stream = torch.xpu.Stream
torch.cuda.default_stream = torch.xpu.current_stream
torch.cuda.current_stream = torch.xpu.current_stream
torch.cuda.stream = torch.xpu.stream
yield
finally:
- # if anything goes wrong, just patch it with a placeholder
- torch.cuda.Event = _EventPlaceholder
+ pass
diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py
index 26c6f8d06bdcd..4d7864e90496a 100644
--- a/vllm/v1/worker/xpu_worker.py
+++ b/vllm/v1/worker/xpu_worker.py
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
+from typing import Any
import torch
import torch.distributed
@@ -37,6 +38,7 @@ class XPUWorker(Worker):
# Torch profiler. Enabled and configured through env vars:
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+ self.profiler: Any | None = None
if envs.VLLM_TORCH_PROFILER_DIR:
torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
@@ -148,7 +150,12 @@ class XPUWorker(Worker):
return int(available_kv_cache_memory)
def init_device(self):
- if self.device_config.device.type == "xpu" and current_platform.is_xpu():
+ device = self.device_config.device
+ if (
+ isinstance(device, torch.device)
+ and device.type == "xpu"
+ and current_platform.is_xpu()
+ ):
self.device = torch.device(f"xpu:{self.local_rank}")
current_platform.set_device(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)