forked from Karylab-cklius/vllm
Compare commits
275
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
902ebda375 | ||
|
|
f8bd8394e3 | ||
|
|
ca81811bfe | ||
|
|
ad8818bb5e | ||
|
|
08e8e99ce7 | ||
|
|
2be765b68a | ||
|
|
16abe6b85a | ||
|
|
1eb61ab34b | ||
|
|
3d962d72ab | ||
|
|
20228cb851 | ||
|
|
7c0d3c5152 | ||
|
|
5b68107411 | ||
|
|
8fb2c135be | ||
|
|
8863c2b25c | ||
|
|
3f72639d36 | ||
|
|
6bc9c8473e | ||
|
|
63ed2409e8 | ||
|
|
95e53d907c | ||
|
|
0346396e94 | ||
|
|
e68b0dad8b | ||
|
|
9cddbdba6d | ||
|
|
49e6b86c91 | ||
|
|
0565f1fdec | ||
|
|
9dbe1fe960 | ||
|
|
a5f89ae296 | ||
|
|
05e8981234 | ||
|
|
899541bdb1 | ||
|
|
d7b2e57097 | ||
|
|
5e034f2e3d | ||
|
|
22970c1626 | ||
|
|
600aaab8d6 | ||
|
|
60446cd684 | ||
|
|
9101dc756c | ||
|
|
025a32f9ed | ||
|
|
19504ac07f | ||
|
|
3df619ac94 | ||
|
|
d74132ca3b | ||
|
|
a34abc49b7 | ||
|
|
d70249e2e9 | ||
|
|
a374532111 | ||
|
|
cee7436a26 | ||
|
|
4c16ba617f | ||
|
|
bde57ab2ed | ||
|
|
9103ed1696 | ||
|
|
46eb30f519 | ||
|
|
0dd63639be | ||
|
|
ef96fa3f1f | ||
|
|
2a4dbe24ea | ||
|
|
8020a60402 | ||
|
|
e15a5ff07b | ||
|
|
6ea001cfb7 | ||
|
|
1c46dea001 | ||
|
|
028599739d | ||
|
|
d1fd802fa3 | ||
|
|
543c23be78 | ||
|
|
b8bf5c45bb | ||
|
|
e6c6f2c79d | ||
|
|
07286ec5a6 | ||
|
|
14fc7a68c7 | ||
|
|
5f2385a4c8 | ||
|
|
a01a1c0d69 | ||
|
|
da6709c9fe | ||
|
|
d83becd503 | ||
|
|
0c9614876e | ||
|
|
583a90e005 | ||
|
|
52d428295d | ||
|
|
c60578de0a | ||
|
|
80fead8bf6 | ||
|
|
e45946bd91 | ||
|
|
ea6d067a2a | ||
|
|
abd9224280 | ||
|
|
4dc0d606b7 | ||
|
|
ac0675ff6b | ||
|
|
e18464a57d | ||
|
|
1963245ed1 | ||
|
|
0308901975 | ||
|
|
aaf4b70aae | ||
|
|
3adffd5b90 | ||
|
|
97ba96fbe9 | ||
|
|
94578127a4 | ||
|
|
2612ba9285 | ||
|
|
1f8b7c536b | ||
|
|
0a0aa07747 | ||
|
|
f9e2a75a1e | ||
|
|
a4d5d663e2 | ||
|
|
657e9c0e18 | ||
|
|
308feab33f | ||
|
|
28ae32a5d3 | ||
|
|
f32c629eb4 | ||
|
|
cd4a95e3aa | ||
|
|
d5ec6c056f | ||
|
|
08d954f036 | ||
|
|
ac9f9330e6 | ||
|
|
2d0c5b630e | ||
|
|
34cd32fe30 | ||
|
|
8e27663b6a | ||
|
|
7cdf7e2fe0 | ||
|
|
bbf80ede43 | ||
|
|
4505849b30 | ||
|
|
db07433ce5 | ||
|
|
e02706d2d2 | ||
|
|
b474782ad7 | ||
|
|
55212c1404 | ||
|
|
e7b68f4d6c | ||
|
|
1a19e9cd87 | ||
|
|
c8ed39b9dd | ||
|
|
020732800c | ||
|
|
dc77cb7129 | ||
|
|
bde38c11df | ||
|
|
707b240d7e | ||
|
|
29ce48221c | ||
|
|
7a05d2dc65 | ||
|
|
a1648c4045 | ||
|
|
e2d49ec2a4 | ||
|
|
8413868dab | ||
|
|
8ff4a99566 | ||
|
|
a4ec0c5595 | ||
|
|
0fa8dd24d2 | ||
|
|
6ebe34d6fa | ||
|
|
11cec296dd | ||
|
|
5825bbc1f7 | ||
|
|
d62cfe546d | ||
|
|
6cdf015c3c | ||
|
|
5d3b6097ad | ||
|
|
e74698c27a | ||
|
|
aa125ecf0e | ||
|
|
f16bfbe5bc | ||
|
|
87e07a6b46 | ||
|
|
7508243249 | ||
|
|
83e1c76dbe | ||
|
|
a563866b48 | ||
|
|
a3d909ad2b | ||
|
|
49568d5cf9 | ||
|
|
b8112c1d85 | ||
|
|
eaba8ece77 | ||
|
|
fe86be66c5 | ||
|
|
1da3a5441a | ||
|
|
72c068b8e0 | ||
|
|
7645bc524b | ||
|
|
1123a87892 | ||
|
|
03fd76c570 | ||
|
|
59d260f5e4 | ||
|
|
18d4e481d0 | ||
|
|
2972a05473 | ||
|
|
5576227bc1 | ||
|
|
d1b6fe007f | ||
|
|
04a49669d1 | ||
|
|
96fcd3c267 | ||
|
|
1f214290d6 | ||
|
|
8cbdc7eb94 | ||
|
|
b634e619bb | ||
|
|
eac3b96ec0 | ||
|
|
573a1d1119 | ||
|
|
33156f56e0 | ||
|
|
107cf8e92f | ||
|
|
63baa28cf5 | ||
|
|
e5173d3bac | ||
|
|
d3235cb503 | ||
|
|
287b37cda4 | ||
|
|
791b2fc30a | ||
|
|
be6a81f31b | ||
|
|
2ab441befe | ||
|
|
9572f74f15 | ||
|
|
5f2a473ff3 | ||
|
|
6b2a672e47 | ||
|
|
f1b1bea5c3 | ||
|
|
cddbc2b4b2 | ||
|
|
087a138963 | ||
|
|
c4041f37a4 | ||
|
|
a79079feef | ||
|
|
9f6dcb71ae | ||
|
|
8dd2419fa9 | ||
|
|
39d82005f7 | ||
|
|
25eef3dc2e | ||
|
|
0d7667419f | ||
|
|
5dcd7ef1f2 | ||
|
|
ffc0a2798b | ||
|
|
10ef65eded | ||
|
|
6170d47d22 | ||
|
|
0ada960a20 | ||
|
|
c907d22158 | ||
|
|
f347ac6c34 | ||
|
|
05f47bd8d2 | ||
|
|
bf184a6621 | ||
|
|
30399cc725 | ||
|
|
b89443b8d9 | ||
|
|
1d9e9ae8a4 | ||
|
|
b7036c87a1 | ||
|
|
cc6dafaef2 | ||
|
|
1ab055efe6 | ||
|
|
b665bbc2d4 | ||
|
|
974138751b | ||
|
|
41cfa50632 | ||
|
|
d111bc53ad | ||
|
|
0790f07695 | ||
|
|
1f33e38e81 | ||
|
|
59fe6f298e | ||
|
|
e7596371a4 | ||
|
|
0dd5dee9b9 | ||
|
|
4614c5a539 | ||
|
|
482914849c | ||
|
|
efeaac92f2 | ||
|
|
55caa6051d | ||
|
|
c7a79d41a0 | ||
|
|
6409004b26 | ||
|
|
aafd4d2354 | ||
|
|
0a2c2dc3f1 | ||
|
|
f09c5feb7c | ||
|
|
1b8af957f6 | ||
|
|
a051525e07 | ||
|
|
5b833be49e | ||
|
|
873480d133 | ||
|
|
6f351548b2 | ||
|
|
364a8bc6dc | ||
|
|
9a1d20a89c | ||
|
|
309a8f66ee | ||
|
|
e5d427e93a | ||
|
|
2a42ae790d | ||
|
|
d49899732e | ||
|
|
dba95378a6 | ||
|
|
ada6f91d56 | ||
|
|
8becf146bd | ||
|
|
c07163663d | ||
|
|
f7008ce1c4 | ||
|
|
4e67a8f616 | ||
|
|
142c4d1738 | ||
|
|
6f5e653383 | ||
|
|
22dffca982 | ||
|
|
4c73be14e0 | ||
|
|
2f4bdee61e | ||
|
|
28c94770ad | ||
|
|
af8fd73051 | ||
|
|
d3e477c013 | ||
|
|
02809af1e7 | ||
|
|
cbd4690a03 | ||
|
|
96860af655 | ||
|
|
0202971a48 | ||
|
|
2c1a4f2488 | ||
|
|
6444824873 | ||
|
|
bf0f3a4638 | ||
|
|
e0327c9db2 | ||
|
|
14df02b4e1 | ||
|
|
6ebb66ccea | ||
|
|
43d384bab4 | ||
|
|
db318326a5 | ||
|
|
799b5721f6 | ||
|
|
97ca4c3b60 | ||
|
|
ee2e69d6cd | ||
|
|
7101e0851f | ||
|
|
e9717801bd | ||
|
|
da71d44410 | ||
|
|
1fb0209bbc | ||
|
|
81323ea221 | ||
|
|
e1cd7a5faf | ||
|
|
a68e703c32 | ||
|
|
cd1245a184 | ||
|
|
ffec815422 | ||
|
|
d386ab1412 | ||
|
|
ccb309a964 | ||
|
|
2f4e6548ef | ||
|
|
3c98c2d21b | ||
|
|
9513029898 | ||
|
|
f6c0009afa | ||
|
|
776ca1e187 | ||
|
|
af9a7ec255 | ||
|
|
276e03b92c | ||
|
|
32f4e4db00 | ||
|
|
ee21291825 | ||
|
|
af1b07b0c5 | ||
|
|
c77a993cc2 | ||
|
|
fdcc5176be | ||
|
|
5708297e4e | ||
|
|
02dbb933cb | ||
|
|
51e38a8e30 | ||
|
|
d8e38d4939 |
@@ -60,6 +60,7 @@ def launch_lm_eval(eval_config, tp_size):
|
||||
f"add_bos_token=true,"
|
||||
f"trust_remote_code={trust_remote_code},"
|
||||
f"max_model_len={max_model_len},"
|
||||
"allow_deprecated_quantization=True,"
|
||||
)
|
||||
|
||||
env_vars = eval_config.get("env_vars", None)
|
||||
|
||||
@@ -2,6 +2,17 @@
|
||||
|
||||
set -euox pipefail
|
||||
|
||||
# To detect ROCm
|
||||
# Check multiple indicators:
|
||||
if [ -e /dev/kfd ] || \
|
||||
[ -d /opt/rocm ] || \
|
||||
command -v rocm-smi &> /dev/null || \
|
||||
[ -n "${ROCM_HOME:-}" ]; then
|
||||
IS_ROCM=1
|
||||
else
|
||||
IS_ROCM=0
|
||||
fi
|
||||
|
||||
if [[ $# -lt 4 ]]; then
|
||||
echo "Usage: .buildkite/scripts/run-multi-node-test.sh WORKING_DIR NUM_NODES NUM_GPUS DOCKER_IMAGE COMMAND1 COMMAND2 ... COMMANDN"
|
||||
exit 1
|
||||
@@ -26,13 +37,18 @@ for command in "${COMMANDS[@]}"; do
|
||||
echo "$command"
|
||||
done
|
||||
|
||||
|
||||
start_network() {
|
||||
docker network create --subnet=192.168.10.0/24 docker-net
|
||||
}
|
||||
|
||||
start_nodes() {
|
||||
for node in $(seq 0 $(($NUM_NODES-1))); do
|
||||
GPU_DEVICES='"device='
|
||||
if [ "$IS_ROCM" -eq 1 ]; then
|
||||
GPU_DEVICES='--device /dev/kfd --device /dev/dri -e HIP_VISIBLE_DEVICES='
|
||||
else
|
||||
GPU_DEVICES='--gpus "device='
|
||||
fi
|
||||
for node_gpu in $(seq 0 $(($NUM_GPUS - 1))); do
|
||||
DEVICE_NUM=$(($node * $NUM_GPUS + $node_gpu))
|
||||
GPU_DEVICES+=$(($DEVICE_NUM))
|
||||
@@ -40,7 +56,9 @@ start_nodes() {
|
||||
GPU_DEVICES+=','
|
||||
fi
|
||||
done
|
||||
GPU_DEVICES+='"'
|
||||
if [ "$IS_ROCM" -eq 0 ]; then
|
||||
GPU_DEVICES+='"'
|
||||
fi
|
||||
|
||||
# start the container in detached mode
|
||||
# things to note:
|
||||
@@ -49,7 +67,7 @@ start_nodes() {
|
||||
# 3. map the huggingface cache directory to the container
|
||||
# 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
|
||||
# starting from 192.168.10.11)
|
||||
docker run -d --gpus "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
|
||||
docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
|
||||
--network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
|
||||
/bin/bash -c "tail -f /dev/null"
|
||||
|
||||
+36
-15
@@ -162,10 +162,7 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/test_vision_embeds.py
|
||||
# Need tf32 to avoid conflicting precision issue with terratorch on ROCm.
|
||||
# TODO: Remove after next torch update
|
||||
- VLLM_FLOAT32_MATMUL_PRECISION="tf32" pytest -v -s entrypoints/openai/test_vision_embeds.py
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration Test (API Server 2)
|
||||
@@ -202,6 +199,21 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration Test (Responses API)
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental]
|
||||
agent_pool: mi325_1
|
||||
# grade: Blocking
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Distributed Tests (4 GPUs) # 35min
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental]
|
||||
@@ -519,8 +531,7 @@ steps:
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
commands:
|
||||
- pytest -v -s samplers
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
|
||||
- pytest -v -s -m 'not skip_v1' samplers
|
||||
|
||||
- label: LoRA Test %N # 20min each
|
||||
timeout_in_minutes: 30
|
||||
@@ -734,7 +745,7 @@ steps:
|
||||
|
||||
- label: Quantization Test # 70min
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental]
|
||||
mirror_hardwares: [amdexperimental, amdproduction]
|
||||
agent_pool: mi325_1
|
||||
# grade: Blocking
|
||||
source_file_dependencies:
|
||||
@@ -989,9 +1000,7 @@ steps:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pip freeze | grep -E 'torch'
|
||||
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing --ignore models/multimodal/pooling/test_prithvi_mae.py
|
||||
# Need tf32 to avoid conflicting precision issue with terratorch on ROCm.
|
||||
# TODO: Remove after next torch update
|
||||
- VLLM_FLOAT32_MATMUL_PRECISION="tf32" pytest -v -s models/multimodal/pooling/test_prithvi_mae.py -m core_model
|
||||
- pytest -v -s models/multimodal/pooling/test_prithvi_mae.py -m core_model
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
|
||||
|
||||
- label: Multi-Modal Accuracy Eval (Small Models) # 5min
|
||||
@@ -1110,8 +1119,8 @@ steps:
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/v1/attention/backends/mla/cutlass_mla.py
|
||||
- vllm/v1/attention/backends/mla/flashinfer_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/cuda.py
|
||||
- vllm/attention/selector.py
|
||||
commands:
|
||||
- nvidia-smi
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
@@ -1356,9 +1365,7 @@ steps:
|
||||
# end platform plugin tests
|
||||
# begin io_processor plugins test, all the code in between uses the prithvi_io_processor plugin
|
||||
- pip install -e ./plugins/prithvi_io_processor_plugin
|
||||
# Need tf32 to avoid conflicting precision issue with terratorch on ROCm.
|
||||
# TODO: Remove after next torch update
|
||||
- VLLM_FLOAT32_MATMUL_PRECISION="tf32" pytest -v -s plugins_tests/test_io_processor_plugins.py
|
||||
- pytest -v -s plugins_tests/test_io_processor_plugins.py
|
||||
- pip uninstall prithvi_io_processor_plugin -y
|
||||
# end io_processor plugins test
|
||||
# begin stat_logger plugins test
|
||||
@@ -1455,7 +1462,21 @@ steps:
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- VLLM_ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh
|
||||
- VLLM_ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP NixlConnector PD accuracy tests (Distributed) # 15min
|
||||
mirror_hardwares: [amdexperimental]
|
||||
agent_pool: mi325_4
|
||||
# grade: Blocking
|
||||
timeout_in_minutes: 15
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- VLLM_ATTENTION_BACKEND=ROCM_ATTN DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
##### multi gpus test #####
|
||||
##### A100 test #####
|
||||
|
||||
@@ -144,7 +144,7 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration Test (API Server 2)
|
||||
@@ -177,6 +177,18 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration Test (Responses API)
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental]
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Distributed Tests (4 GPUs) # 35min
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental]
|
||||
@@ -383,7 +395,6 @@ steps:
|
||||
- label: V1 Test attention (B200) # 10min
|
||||
timeout_in_minutes: 30
|
||||
gpu: b200
|
||||
optional: true # TODO(mgoin): disable optional once runner is fixed
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention
|
||||
- tests/v1/attention
|
||||
@@ -944,7 +955,6 @@ steps:
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
optional: true # TODO(mgoin): disable optional once runner is fixed
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- csrc/attention/mla/
|
||||
@@ -956,8 +966,8 @@ steps:
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/v1/attention/backends/mla/cutlass_mla.py
|
||||
- vllm/v1/attention/backends/mla/flashinfer_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/cuda.py
|
||||
- vllm/attention/selector.py
|
||||
commands:
|
||||
- nvidia-smi
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
@@ -986,7 +996,6 @@ steps:
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
optional: true # TODO(mgoin): disable optional once runner is fixed
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
@@ -1054,7 +1063,6 @@ steps:
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
optional: true # TODO(mgoin): disable optional once runner is fixed
|
||||
source_file_dependencies:
|
||||
- tests/quantization/test_blackwell_moe.py
|
||||
- vllm/model_executor/models/deepseek_v2.py
|
||||
@@ -1108,6 +1116,7 @@ steps:
|
||||
- vllm/model_executor/models/
|
||||
- tests/distributed/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- .buildkite/scripts/run-multi-node-test.sh
|
||||
commands:
|
||||
- # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)
|
||||
- VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
@@ -1270,8 +1279,8 @@ steps:
|
||||
commands:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
|
||||
|
||||
- label: NixlConnector PD accuracy tests (Distributed) # 30min
|
||||
timeout_in_minutes: 30
|
||||
- label: NixlConnector PD accuracy tests (Distributed) # 40min
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
source_file_dependencies:
|
||||
@@ -1281,8 +1290,8 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP NixlConnector PD accuracy tests (Distributed)
|
||||
timeout_in_minutes: 30
|
||||
- label: DP EP NixlConnector PD accuracy tests (Distributed) # 15min
|
||||
timeout_in_minutes: 15
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
source_file_dependencies:
|
||||
@@ -1351,6 +1360,14 @@ steps:
|
||||
- CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
|
||||
- label: LM Eval Large Models (H200) # optional
|
||||
timeout_in_minutes: 60
|
||||
gpu: h200
|
||||
optional: true
|
||||
num_gpus: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
|
||||
|
||||
##### B200 test #####
|
||||
- label: Distributed Tests (B200) # optional
|
||||
gpu: b200
|
||||
@@ -1402,3 +1419,26 @@ steps:
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
|
||||
|
||||
##### MoE Refactor (Temporary) Tests #####
|
||||
|
||||
- label: MoE Refactor Integration Test (H100 - TEMPORARY) # optional
|
||||
gpu: h100
|
||||
optional: true
|
||||
num_gpus: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 - TEMPORARY) # optional
|
||||
gpu: b200
|
||||
optional: true
|
||||
num_gpus: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY) # optional
|
||||
gpu: b200
|
||||
optional: true
|
||||
num_gpus: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
|
||||
|
||||
@@ -182,7 +182,7 @@ steps:
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs))
|
||||
timeout_in_minutes: 60
|
||||
|
||||
@@ -34,10 +34,9 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server 2)
|
||||
timeout_in_minutes: 130
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -64,6 +63,14 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Entrypoints V1
|
||||
timeout_in_minutes: 50
|
||||
|
||||
@@ -90,8 +90,8 @@ steps:
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/v1/attention/backends/mla/cutlass_mla.py
|
||||
- vllm/v1/attention/backends/mla/flashinfer_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/cuda.py
|
||||
- vllm/attention/selector.py
|
||||
commands:
|
||||
- nvidia-smi
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
|
||||
+5
-5
@@ -3,7 +3,6 @@
|
||||
|
||||
# This lists cover the "core" components of vLLM that require careful review
|
||||
/vllm/attention @LucasWilkinson
|
||||
/vllm/attention/backends/abstract.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
|
||||
/vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @njhill @22quinn
|
||||
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
|
||||
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
|
||||
@@ -27,6 +26,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
|
||||
|
||||
# vLLM V1
|
||||
/vllm/v1/attention @LucasWilkinson
|
||||
/vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
|
||||
/vllm/v1/attention/backends/mla @pavanimajety
|
||||
/vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety
|
||||
/vllm/v1/attention/backends/triton_attn.py @tdoublep
|
||||
@@ -117,15 +117,15 @@ mkdocs.yaml @hmellor
|
||||
/vllm/transformers_utils/tokenizers/mistral.py @patrickvonplaten
|
||||
|
||||
# Kernels
|
||||
/vllm/attention/ops/chunked_prefill_paged_decode.py @tdoublep
|
||||
/vllm/attention/ops/triton_unified_attention.py @tdoublep
|
||||
/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
|
||||
/vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
|
||||
|
||||
# ROCm related: specify owner with write access to notify AMD folks for careful code review
|
||||
/vllm/**/*rocm* @tjtanaa
|
||||
/docker/Dockerfile.rocm* @gshtras @tjtanaa
|
||||
/vllm/v1/attention/backends/rocm*.py @gshtras @tjtanaa
|
||||
/vllm/v1/attention/backends/mla/rocm*.py @gshtras @tjtanaa
|
||||
/vllm/attention/ops/rocm*.py @gshtras @tjtanaa
|
||||
/vllm/v1/attention/ops/rocm*.py @gshtras @tjtanaa
|
||||
/vllm/model_executor/layers/fused_moe/rocm*.py @gshtras @tjtanaa
|
||||
/csrc/rocm @gshtras @tjtanaa
|
||||
/requirements/*rocm* @tjtanaa
|
||||
@@ -153,7 +153,7 @@ mkdocs.yaml @hmellor
|
||||
/vllm/entrypoints/pooling @noooop
|
||||
/vllm/config/pooler.py @noooop
|
||||
/vllm/pooling_params.py @noooop
|
||||
/vllm/model_executor/layers/pooler.py @noooop
|
||||
/vllm/model_executor/layers/pooler @noooop
|
||||
|
||||
# Security guide and policies
|
||||
/docs/usage/security.md @russellb
|
||||
|
||||
+2
-2
@@ -222,10 +222,10 @@ pull_request_rules:
|
||||
- files~=^csrc/rocm/
|
||||
- files~=^docker/Dockerfile.rocm
|
||||
- files~=^requirements/rocm.*\.txt
|
||||
- files~=^vllm/attention/backends/rocm.*\.py
|
||||
- files~=^vllm/attention/ops/rocm.*\.py
|
||||
- files~=^vllm/model_executor/layers/fused_moe/rocm.*\.py
|
||||
- files~=^vllm/v1/attention/backends/rocm.*\.py
|
||||
- files~=^vllm/v1/attention/backends/mla/rocm.*\.py
|
||||
- files~=^vllm/v1/attention/ops/rocm.*\.py
|
||||
- files~=^tests/kernels/.*_rocm.*\.py
|
||||
- files=vllm/platforms/rocm.py
|
||||
- title~=(?i)AMD
|
||||
|
||||
@@ -227,3 +227,8 @@ ep_kernels_workspace/
|
||||
|
||||
# Allow tracked library source folders under submodules (e.g., benchmarks/lib)
|
||||
!vllm/benchmarks/lib/
|
||||
|
||||
# Generated gRPC protobuf files (compiled at build time from vllm_engine.proto)
|
||||
vllm/grpc/vllm_engine_pb2.py
|
||||
vllm/grpc/vllm_engine_pb2_grpc.py
|
||||
vllm/grpc/vllm_engine_pb2.pyi
|
||||
|
||||
@@ -282,6 +282,7 @@ endif()
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
|
||||
"csrc/cache_kernels.cu"
|
||||
"csrc/cache_kernels_fused.cu"
|
||||
"csrc/attention/paged_attention_v1.cu"
|
||||
"csrc/attention/paged_attention_v2.cu"
|
||||
"csrc/attention/merge_attn_states.cu"
|
||||
|
||||
+15
-32
@@ -1,47 +1,30 @@
|
||||
# Releasing vLLM
|
||||
|
||||
vLLM releases offer a reliable version of the code base, packaged into a binary format that can be conveniently accessed via PyPI. These releases also serve as key milestones for the development team to communicate with the community about newly available features, improvements, and upcoming changes that could affect users, including potential breaking changes.
|
||||
vLLM releases offer a reliable version of the code base, packaged into a binary format that can be conveniently accessed via [PyPI](https://pypi.org/project/vllm). These releases also serve as key milestones for the development team to communicate with the community about newly available features, improvements, and upcoming changes that could affect users, including potential breaking changes.
|
||||
|
||||
## Release Versioning
|
||||
## Release Cadence and Versioning
|
||||
|
||||
vLLM uses a “right-shifted” versioning scheme where a new patch release is out every 2 weeks. And patch releases contain features and bug fixes (as opposed to semver where patch release contains only backwards-compatible bug fixes). When critical fixes need to be made, special release post1 is released.
|
||||
We aim to have a regular release every 2 weeks. Since v0.12.0, regular releases increment the minor version rather than patch version. The list of past releases can be found [here](https://vllm.ai/releases).
|
||||
|
||||
* _major_ major architectural milestone and when incompatible API changes are made, similar to PyTorch 2.0.
|
||||
* _minor_ major features
|
||||
* _patch_ features and backwards-compatible bug fixes
|
||||
* _post1_ or _patch-1_ backwards-compatible bug fixes, either explicit or implicit post release
|
||||
Our version numbers are expressed in the form `vX.Y.Z`, where `X` is the major version, `Y` is the minor version, and `Z` is the patch version. They are incremented according to the following rules:
|
||||
|
||||
## Release Cadence
|
||||
* _Major_ releases are reserved for architectural milestones involving sweeping API changes, similar to PyTorch 2.0.
|
||||
* _Minor_ releases correspond to regular releases, which include new features, bug fixes and other backwards-compatible changes.
|
||||
* _Patch_ releases correspond to special releases for new models, as well as emergency patches for critical performance, functionality and security issues.
|
||||
|
||||
Patch release is released on bi-weekly basis. Post release 1-3 days after patch release and uses same branch as patch release.
|
||||
Following is the release cadence for year 2025. All future release dates below are tentative. Please note: Post releases are optional.
|
||||
This versioning scheme is similar to [SemVer](https://semver.org/) for compatibility purposes, except that backwards compatibility is only guaranteed for a limited number of minor releases (see our [deprecation policy](https://docs.vllm.ai/en/latest/contributing/deprecation_policy) for details).
|
||||
|
||||
| Release Date | Patch release versions | Post Release versions |
|
||||
| --- | --- | --- |
|
||||
| Jan 2025 | 0.7.0 | --- |
|
||||
| Feb 2025 | 0.7.1, 0.7.2, 0.7.3 | --- |
|
||||
| Mar 2025 | 0.7.4, 0.7.5 | --- |
|
||||
| Apr 2025 | 0.7.6, 0.7.7 | --- |
|
||||
| May 2025 | 0.7.8, 0.7.9 | --- |
|
||||
| Jun 2025 | 0.7.10, 0.7.11 | --- |
|
||||
| Jul 2025 | 0.7.12, 0.7.13 | --- |
|
||||
| Aug 2025 | 0.7.14, 0.7.15 | --- |
|
||||
| Sep 2025 | 0.7.16, 0.7.17 | --- |
|
||||
| Oct 2025 | 0.7.18, 0.7.19 | --- |
|
||||
| Nov 2025 | 0.7.20, 0.7.21 | --- |
|
||||
| Dec 2025 | 0.7.22, 0.7.23 | --- |
|
||||
|
||||
## Release branch
|
||||
## Release Branch
|
||||
|
||||
Each release is built from a dedicated release branch.
|
||||
|
||||
* For _major_, _minor_, _patch_ releases, the release branch cut is performed 1-2 days before release is live.
|
||||
* For post releases, previously cut release branch is reused
|
||||
* Release builds are triggered via push to RC tag like vX.Y.Z-rc1 . This enables us to build and test multiple RCs for each release.
|
||||
* Final tag : vX.Y.Z does not trigger the build but used for Release notes and assets.
|
||||
* After branch cut is created we monitor the main branch for any reverts and apply these reverts to a release branch.
|
||||
* For _major_ and _minor_ releases, the release branch cut is performed 1-2 days before release is live.
|
||||
* For _patch_ releases, previously cut release branch is reused.
|
||||
* Release builds are triggered via push to RC tag like `vX.Y.Z-rc1`. This enables us to build and test multiple RCs for each release.
|
||||
* Final tag: `vX.Y.Z` does not trigger the build but used for Release notes and assets.
|
||||
* After branch cut is created, we monitor the main branch for any reverts and apply these reverts to a release branch.
|
||||
|
||||
## Release Cherry-Pick Criteria
|
||||
### Cherry-Pick Criteria
|
||||
|
||||
After branch cut, we approach finalizing the release branch with clear criteria on what cherry picks are allowed in. Note: a cherry pick is a process to land a PR in the release branch after branch cut. These are typically limited to ensure that the team has sufficient time to complete a thorough round of testing on a stable code base.
|
||||
|
||||
|
||||
@@ -135,7 +135,6 @@ def benchmark_batched_propose(args):
|
||||
block_sizes=[16],
|
||||
)
|
||||
dummy_input_batch._req_ids = list(str(id) for id in range(args.num_req))
|
||||
dummy_input_batch.spec_decode_unsupported_reqs = ()
|
||||
dummy_input_batch.num_tokens_no_spec = [args.num_token] * args.num_req
|
||||
dummy_input_batch.token_ids_cpu = np.random.randint(
|
||||
0, 20, (args.num_req, args.num_token)
|
||||
@@ -151,10 +150,8 @@ def benchmark_batched_propose(args):
|
||||
start = time.time()
|
||||
runner.drafter.propose(
|
||||
sampled_token_ids,
|
||||
dummy_input_batch.req_ids,
|
||||
dummy_input_batch.num_tokens_no_spec,
|
||||
dummy_input_batch.token_ids_cpu,
|
||||
dummy_input_batch.spec_decode_unsupported_reqs,
|
||||
)
|
||||
end = time.time()
|
||||
print(f"Iteration time (s): {end - start}")
|
||||
|
||||
@@ -343,7 +343,9 @@ def bench(
|
||||
return bench_int8(dtype, m, k, n, label, sub_label)
|
||||
if dtype == torch.float8_e4m3fn:
|
||||
return bench_fp8(dtype, m, k, n, label, sub_label)
|
||||
raise ValueError("unsupported type")
|
||||
raise ValueError(
|
||||
f"Unsupported dtype {dtype}: should be one of torch.int8, torch.float8_e4m3fn."
|
||||
)
|
||||
|
||||
|
||||
# runner
|
||||
|
||||
@@ -8,10 +8,9 @@ import torch
|
||||
|
||||
import vllm.model_executor.layers.activation # noqa F401
|
||||
from vllm.model_executor.custom_op import CustomOp
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
|
||||
|
||||
batch_size_range = [1, 16, 128]
|
||||
seq_len_range = [1, 16, 64, 1024, 4096]
|
||||
@@ -30,7 +29,7 @@ def benchmark_activation(
|
||||
device = "cuda"
|
||||
num_tokens = batch_size * seq_len
|
||||
dim = intermediate_size
|
||||
current_platform.seed_everything(42)
|
||||
set_random_seed(42)
|
||||
torch.set_default_device(device)
|
||||
|
||||
if func_name == "gelu_and_mul":
|
||||
|
||||
@@ -6,15 +6,19 @@ kernel. Both kernels take in fp8 quantized weights and 16-bit activations,
|
||||
but use different quantization strategies and backends.
|
||||
"""
|
||||
|
||||
import nvtx
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
||||
MoEPrepareAndFinalizeNoEP,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
# Weight shapes for different models: [num_experts, topk, hidden_size,
|
||||
# intermediate_size]
|
||||
@@ -58,6 +62,7 @@ def bench_run(
|
||||
per_out_ch: bool,
|
||||
mkn: tuple[int, int, int],
|
||||
):
|
||||
init_workspace_manager(torch.cuda.current_device())
|
||||
(m, k, n) = mkn
|
||||
|
||||
dtype = torch.half
|
||||
@@ -120,85 +125,6 @@ def bench_run(
|
||||
# Force per-tensor quantization for all cases
|
||||
per_act_token = False
|
||||
|
||||
# Create stride tensors for CUTLASS
|
||||
ab_strides1 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
|
||||
ab_strides2 = torch.full((num_experts,), n, dtype=torch.int64, device=device)
|
||||
c_strides1 = torch.full((num_experts,), 2 * n, dtype=torch.int64, device=device)
|
||||
c_strides2 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
|
||||
|
||||
def run_triton_moe(
|
||||
a: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
w1_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
a1_scale: torch.Tensor,
|
||||
a2_scale: torch.Tensor,
|
||||
num_repeats: int,
|
||||
):
|
||||
quant_config = fp8_w8a8_moe_quant_config(
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
per_act_token_quant=per_act_token,
|
||||
per_out_ch_quant=per_out_ch,
|
||||
)
|
||||
|
||||
for _ in range(num_repeats):
|
||||
fused_experts(
|
||||
a,
|
||||
w1,
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
def run_cutlass_moe_fp8(
|
||||
a: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
ab_strides1: torch.Tensor,
|
||||
ab_strides2: torch.Tensor,
|
||||
c_strides1: torch.Tensor,
|
||||
c_strides2: torch.Tensor,
|
||||
w1_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
a1_scale: torch.Tensor,
|
||||
a2_scale: torch.Tensor,
|
||||
num_repeats: int,
|
||||
):
|
||||
quant_config = fp8_w8a8_moe_quant_config(
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
per_act_token_quant=per_act_token,
|
||||
per_out_ch_quant=per_out_ch,
|
||||
)
|
||||
|
||||
for _ in range(num_repeats):
|
||||
with nvtx.annotate("cutlass_moe_fp8", color="blue"):
|
||||
cutlass_moe_fp8(
|
||||
a=a,
|
||||
w1_q=w1,
|
||||
w2_q=w2,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
ab_strides1=ab_strides1,
|
||||
ab_strides2=ab_strides2,
|
||||
c_strides1=c_strides1,
|
||||
c_strides2=c_strides2,
|
||||
quant_config=quant_config,
|
||||
activation="silu",
|
||||
global_num_experts=num_experts,
|
||||
)
|
||||
|
||||
# Pre-create quantization config to avoid creating it inside CUDA graph
|
||||
quant_config = fp8_w8a8_moe_quant_config(
|
||||
w1_scale=w1_scale,
|
||||
@@ -209,23 +135,30 @@ def bench_run(
|
||||
per_out_ch_quant=per_out_ch,
|
||||
)
|
||||
|
||||
fn = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
CutlassExpertsFp8(
|
||||
out_dtype=a.dtype,
|
||||
e=num_experts,
|
||||
n=n,
|
||||
k=k,
|
||||
quant_config=quant_config,
|
||||
device=w1.device,
|
||||
),
|
||||
)
|
||||
|
||||
# Create CUDA graphs for CUTLASS (match benchmark_moe.py pattern exactly)
|
||||
cutlass_stream = torch.cuda.Stream()
|
||||
cutlass_graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):
|
||||
# Capture 10 invocations like benchmark_moe.py
|
||||
for _ in range(10):
|
||||
cutlass_moe_fp8(
|
||||
a=a,
|
||||
w1_q=w1_fp8q_cutlass,
|
||||
w2_q=w2_fp8q_cutlass,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
ab_strides1=ab_strides1,
|
||||
ab_strides2=ab_strides2,
|
||||
c_strides1=c_strides1,
|
||||
c_strides2=c_strides2,
|
||||
quant_config=quant_config,
|
||||
fn(
|
||||
a,
|
||||
w1_fp8q_cutlass,
|
||||
w2_fp8q_cutlass,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation="silu",
|
||||
global_num_experts=num_experts,
|
||||
)
|
||||
@@ -297,6 +230,10 @@ def bench_run(
|
||||
|
||||
|
||||
def main(args):
|
||||
# Initialize workspace manager (required for CUTLASS MoE kernels)
|
||||
device = torch.device("cuda:0")
|
||||
init_workspace_manager(device)
|
||||
|
||||
print("Benchmarking models:")
|
||||
for i, model in enumerate(args.models):
|
||||
print(f"[{i}] {model}")
|
||||
|
||||
+39
-19
@@ -11,16 +11,23 @@ import nvtx
|
||||
import torch
|
||||
import torch.utils.benchmark as benchmark
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
fp8_w8a8_moe_quant_config,
|
||||
nvfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
|
||||
CutlassExpertsFp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
||||
MoEPrepareAndFinalizeNoEP,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
WEIGHT_SHAPES_MOE = {
|
||||
"nvidia/DeepSeek-R1-FP4": [
|
||||
@@ -187,19 +194,24 @@ def bench_run(
|
||||
g1_alphas=w1_gs,
|
||||
g2_alphas=w2_gs,
|
||||
)
|
||||
|
||||
kernel = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
|
||||
CutlassExpertsFp4(
|
||||
out_dtype=dtype,
|
||||
max_experts_per_worker=e,
|
||||
quant_config=quant_config,
|
||||
),
|
||||
)
|
||||
|
||||
for _ in range(num_repeats):
|
||||
with nvtx.annotate("cutlass_moe_fp4", color="green"):
|
||||
cutlass_moe_fp4(
|
||||
a=a,
|
||||
w1_fp4=w1_fp4,
|
||||
w2_fp4=w2_fp4,
|
||||
kernel(
|
||||
hidden_states=a,
|
||||
w1=w1_fp4,
|
||||
w2=w2_fp4,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
m=m,
|
||||
n=n,
|
||||
k=k,
|
||||
e=num_experts,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
def run_cutlass_from_graph(
|
||||
@@ -229,20 +241,24 @@ def bench_run(
|
||||
g2_alphas=w2_gs,
|
||||
)
|
||||
|
||||
kernel = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
|
||||
CutlassExpertsFp4(
|
||||
out_dtype=dtype,
|
||||
max_experts_per_worker=e,
|
||||
quant_config=quant_config,
|
||||
),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(
|
||||
VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
|
||||
):
|
||||
return cutlass_moe_fp4(
|
||||
a=a,
|
||||
w1_fp4=w1_fp4,
|
||||
w2_fp4=w2_fp4,
|
||||
return kernel(
|
||||
hidden_states=a,
|
||||
w1=w1_fp4,
|
||||
w2=w2_fp4,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
m=m,
|
||||
n=n,
|
||||
k=k,
|
||||
e=num_experts,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
def run_triton_from_graph(
|
||||
@@ -441,6 +457,10 @@ def bench_run(
|
||||
|
||||
|
||||
def main(args):
|
||||
# Initialize workspace manager (required for CUTLASS MoE kernels)
|
||||
device = torch.device("cuda:0")
|
||||
init_workspace_manager(device)
|
||||
|
||||
print("Benchmarking models:")
|
||||
for i, model in enumerate(args.models):
|
||||
print(f"[{i}] {model}")
|
||||
@@ -5,15 +5,20 @@ import torch
|
||||
import torch.utils.benchmark as benchmark
|
||||
from benchmark_shapes import WEIGHT_SHAPES_MOE
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import (
|
||||
fused_experts,
|
||||
fused_topk,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
||||
MoEPrepareAndFinalizeNoEP,
|
||||
)
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
DEFAULT_MODELS = [
|
||||
"mistralai/Mixtral-8x7B-Instruct-v0.1",
|
||||
@@ -44,6 +49,7 @@ def bench_run(
|
||||
per_out_ch: bool,
|
||||
mkn: tuple[int, int, int],
|
||||
):
|
||||
init_workspace_manager(torch.cuda.current_device())
|
||||
label = "Quant Matmul"
|
||||
|
||||
sub_label = (
|
||||
@@ -81,11 +87,6 @@ def bench_run(
|
||||
a, score, topk, renormalize=False
|
||||
)
|
||||
|
||||
ab_strides1 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
|
||||
ab_strides2 = torch.full((num_experts,), n, device="cuda", dtype=torch.int64)
|
||||
c_strides1 = torch.full((num_experts,), 2 * n, device="cuda", dtype=torch.int64)
|
||||
c_strides2 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
|
||||
|
||||
def run_triton_moe(
|
||||
a: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
@@ -119,10 +120,6 @@ def bench_run(
|
||||
w2: torch.Tensor,
|
||||
w1_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
ab_strides1: torch.Tensor,
|
||||
ab_strides2: torch.Tensor,
|
||||
c_strides1: torch.Tensor,
|
||||
c_strides2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
per_act_token: bool,
|
||||
@@ -134,31 +131,29 @@ def bench_run(
|
||||
per_act_token_quant=per_act_token,
|
||||
)
|
||||
|
||||
for _ in range(num_repeats):
|
||||
cutlass_moe_fp8(
|
||||
a,
|
||||
w1,
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
ab_strides1,
|
||||
ab_strides2,
|
||||
c_strides1,
|
||||
c_strides2,
|
||||
fn = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
CutlassExpertsFp8(
|
||||
out_dtype=a.dtype,
|
||||
# NOTE(rob): w2 is shaped as [E, hidden, intermediate]
|
||||
e=w2.shape[0],
|
||||
n=w2.shape[2],
|
||||
k=w2.shape[1],
|
||||
quant_config=quant_config,
|
||||
)
|
||||
device=w1.device,
|
||||
),
|
||||
)
|
||||
|
||||
for _ in range(num_repeats):
|
||||
fn(a, w1, w2, topk_weights, topk_ids)
|
||||
|
||||
def run_cutlass_from_graph(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
w1_q: torch.Tensor,
|
||||
w2_q: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w1_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
ab_strides1: torch.Tensor,
|
||||
ab_strides2: torch.Tensor,
|
||||
c_strides1: torch.Tensor,
|
||||
c_strides2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
):
|
||||
@@ -168,21 +163,23 @@ def bench_run(
|
||||
per_act_token_quant=per_act_token,
|
||||
)
|
||||
|
||||
fn = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
CutlassExpertsFp8(
|
||||
out_dtype=a.dtype,
|
||||
# NOTE(rob): w2 is shaped as [E, hidden, intermediate]
|
||||
e=w2.shape[0],
|
||||
n=w2.shape[2],
|
||||
k=w2.shape[1],
|
||||
quant_config=quant_config,
|
||||
device=w1.device,
|
||||
),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(
|
||||
VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
|
||||
):
|
||||
return cutlass_moe_fp8(
|
||||
a,
|
||||
w1_q,
|
||||
w2_q,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
ab_strides1,
|
||||
ab_strides2,
|
||||
c_strides1,
|
||||
c_strides2,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
return fn(a, w1, w2, topk_weights, topk_ids)
|
||||
|
||||
def run_triton_from_graph(
|
||||
a: torch.Tensor,
|
||||
@@ -226,10 +223,6 @@ def bench_run(
|
||||
w2_q,
|
||||
w1_scale,
|
||||
w2_scale,
|
||||
ab_strides1,
|
||||
ab_strides2,
|
||||
c_strides1,
|
||||
c_strides2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
)
|
||||
@@ -267,10 +260,6 @@ def bench_run(
|
||||
"w1_scale": w1_scale,
|
||||
"w2_scale": w2_scale,
|
||||
"per_act_token": per_act_token,
|
||||
"ab_strides1": ab_strides1,
|
||||
"ab_strides2": ab_strides2,
|
||||
"c_strides1": c_strides1,
|
||||
"c_strides2": c_strides2,
|
||||
# cuda graph params
|
||||
"cutlass_graph": cutlass_graph,
|
||||
"triton_graph": triton_graph,
|
||||
@@ -329,10 +318,6 @@ def bench_run(
|
||||
w2_q,
|
||||
w1_scale,
|
||||
w2_scale,
|
||||
ab_strides1,
|
||||
ab_strides2,
|
||||
c_strides1,
|
||||
c_strides2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
per_act_token,
|
||||
@@ -341,7 +326,7 @@ def bench_run(
|
||||
|
||||
results.append(
|
||||
benchmark.Timer(
|
||||
stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, ab_strides1, ab_strides2, c_strides1, c_strides2, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
|
||||
stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
|
||||
globals=globals,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
@@ -364,6 +349,10 @@ def bench_run(
|
||||
|
||||
|
||||
def main(args):
|
||||
# Initialize workspace manager (required for CUTLASS MoE kernels)
|
||||
device = torch.device("cuda:0")
|
||||
init_workspace_manager(device)
|
||||
|
||||
print("Benchmarking models:")
|
||||
for i, model in enumerate(args.models):
|
||||
print(f"[{i}] {model}")
|
||||
|
||||
@@ -6,9 +6,8 @@ import time
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
@@ -22,7 +21,7 @@ def main(
|
||||
num_warmup_iters: int = 5,
|
||||
num_iters: int = 100,
|
||||
) -> None:
|
||||
current_platform.seed_everything(seed)
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
layer = RMSNorm(hidden_size).to(dtype=dtype)
|
||||
|
||||
@@ -24,6 +24,7 @@ from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.config import get_config
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
@@ -47,8 +48,6 @@ def clear_triton_cache():
|
||||
|
||||
# Try to clear Triton's runtime cache
|
||||
try:
|
||||
import triton
|
||||
|
||||
if (
|
||||
hasattr(triton, "runtime")
|
||||
and hasattr(triton.runtime, "cache")
|
||||
@@ -431,7 +430,7 @@ def merge_unique_dicts(list1, list2):
|
||||
class BenchmarkWorker:
|
||||
def __init__(self, seed: int) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
current_platform.seed_everything(seed)
|
||||
set_random_seed(seed)
|
||||
self.seed = seed
|
||||
# Get the device ID to allocate tensors and kernels
|
||||
# on the respective GPU. This is required for Ray to work
|
||||
@@ -451,7 +450,7 @@ class BenchmarkWorker:
|
||||
block_quant_shape: list[int] = None,
|
||||
use_deep_gemm: bool = False,
|
||||
) -> tuple[dict[str, int], float]:
|
||||
current_platform.seed_everything(self.seed)
|
||||
set_random_seed(self.seed)
|
||||
dtype_str = _get_config_dtype_str(
|
||||
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
|
||||
)
|
||||
|
||||
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
|
||||
from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
@@ -261,7 +262,7 @@ def benchmark_unpermute(
|
||||
class BenchmarkWorker:
|
||||
def __init__(self, seed: int) -> None:
|
||||
torch.set_default_device("cuda")
|
||||
current_platform.seed_everything(seed)
|
||||
set_random_seed(seed)
|
||||
self.seed = seed
|
||||
# Get the device ID to allocate tensors and kernels
|
||||
# on the respective GPU. This is required for Ray to work
|
||||
@@ -279,7 +280,7 @@ class BenchmarkWorker:
|
||||
use_int8_w8a16: bool,
|
||||
use_customized_permute: bool = False,
|
||||
) -> tuple[dict[str, int], float]:
|
||||
current_platform.seed_everything(self.seed)
|
||||
set_random_seed(self.seed)
|
||||
|
||||
permute_time = benchmark_permute(
|
||||
num_tokens,
|
||||
|
||||
@@ -37,9 +37,9 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.config import get_config
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
@@ -94,7 +94,7 @@ def benchmark_mrope(
|
||||
benchmark_iter: int = 100,
|
||||
csv_writer=None,
|
||||
):
|
||||
current_platform.seed_everything(seed)
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
# the parameters to compute the q k v size based on tp_size
|
||||
mrope_helper_class = get_rope(
|
||||
|
||||
@@ -13,6 +13,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import (
|
||||
STR_DTYPE_TO_TORCH_DTYPE,
|
||||
create_kv_caches_with_random,
|
||||
set_random_seed,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -38,7 +39,7 @@ def main(
|
||||
device: str = "cuda",
|
||||
kv_cache_dtype: str | None = None,
|
||||
) -> None:
|
||||
current_platform.seed_everything(seed)
|
||||
set_random_seed(seed)
|
||||
|
||||
scale = float(1.0 / (head_size**0.5))
|
||||
query = torch.empty(
|
||||
|
||||
@@ -6,9 +6,8 @@ import time
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
@@ -23,7 +22,7 @@ def main(
|
||||
num_warmup_iters: int = 5,
|
||||
num_iters: int = 100,
|
||||
) -> None:
|
||||
current_platform.seed_everything(seed)
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
|
||||
|
||||
@@ -8,11 +8,11 @@ from tabulate import tabulate
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import (
|
||||
STR_DTYPE_TO_TORCH_DTYPE,
|
||||
create_kv_caches_with_random,
|
||||
set_random_seed,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -36,7 +36,7 @@ def run_benchmark(
|
||||
if kv_cache_dtype == "fp8" and head_size % 16:
|
||||
raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")
|
||||
|
||||
current_platform.seed_everything(42)
|
||||
set_random_seed(42)
|
||||
torch.set_default_device(device)
|
||||
|
||||
# create random key / value tensors [T, H, D].
|
||||
|
||||
@@ -7,15 +7,15 @@ import torch
|
||||
from tabulate import tabulate
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.attention.ops.triton_reshape_and_cache_flash import (
|
||||
triton_reshape_and_cache_flash,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import (
|
||||
STR_DTYPE_TO_TORCH_DTYPE,
|
||||
create_kv_caches_with_random_flash,
|
||||
set_random_seed,
|
||||
)
|
||||
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
|
||||
triton_reshape_and_cache_flash,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -49,7 +49,7 @@ def run_benchmark(
|
||||
if implementation == "triton" and kv_cache_layout == "HND":
|
||||
return float("nan") # Triton does not support HND layout yet.
|
||||
|
||||
current_platform.seed_everything(42)
|
||||
set_random_seed(42)
|
||||
torch.set_default_device(device)
|
||||
|
||||
# create random key / value tensors [T, H, D].
|
||||
|
||||
@@ -23,9 +23,9 @@ import torch
|
||||
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
|
||||
persistent_masked_m_silu_mul_quant,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@triton.jit
|
||||
@@ -207,7 +207,7 @@ def benchmark(
|
||||
):
|
||||
def generate_data(seed_offset=0):
|
||||
"""Generate input data with given seed offset"""
|
||||
current_platform.seed_everything(42 + seed_offset)
|
||||
set_random_seed(42 + seed_offset)
|
||||
y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous()
|
||||
|
||||
if gen_strategy == "random_imbalanced":
|
||||
|
||||
@@ -0,0 +1,272 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import functools
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm._custom_ops import (
|
||||
cpu_attention_with_kv_cache,
|
||||
cpu_attn_get_scheduler_metadata,
|
||||
cpu_attn_reshape_and_cache,
|
||||
)
|
||||
from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
|
||||
from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa
|
||||
|
||||
|
||||
def get_attn_isa(
|
||||
block_size: int | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
if block_size and dtype:
|
||||
return _get_attn_isa(dtype, block_size)
|
||||
else:
|
||||
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
|
||||
return "neon"
|
||||
elif torch._C._cpu._is_amx_tile_supported():
|
||||
return "amx"
|
||||
else:
|
||||
return "vec"
|
||||
|
||||
|
||||
# rand number generation takes too much time, cache rand tensors
|
||||
@functools.lru_cache(maxsize=128, typed=False)
|
||||
def tensor_cache(
|
||||
elem_num: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
tensor = torch.randn(elem_num, dtype=dtype)
|
||||
return tensor
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int = None,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
block_size: int = 128,
|
||||
num_blocks: int = 4096,
|
||||
use_sink: bool = False,
|
||||
enable_kv_split: bool = False,
|
||||
isa: str | None = None,
|
||||
seed: int = 0,
|
||||
iters: int = 20,
|
||||
) -> None:
|
||||
current_platform.seed_everything(seed)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
token_num = sum(query_lens)
|
||||
|
||||
if isa is None:
|
||||
isa = get_attn_isa(block_size, dtype)
|
||||
|
||||
s_aux = (
|
||||
15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
|
||||
)
|
||||
|
||||
query = tensor_cache(
|
||||
elem_num=token_num * num_query_heads * head_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
query = query.view(
|
||||
token_num,
|
||||
num_query_heads,
|
||||
head_size,
|
||||
)
|
||||
|
||||
key_value = tensor_cache(
|
||||
elem_num=2 * num_blocks * num_kv_heads * block_size * head_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
key_value = key_value.view(
|
||||
2,
|
||||
num_blocks,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
)
|
||||
key_cache, value_cache = key_value.unbind(0)
|
||||
|
||||
# KV cache for CPU attention
|
||||
packed_key_cache = torch.empty(
|
||||
num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
|
||||
)
|
||||
packed_value_cache = torch.empty_like(packed_key_cache)
|
||||
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
# use reshape_and_cache to pack key_cache and value_cache
|
||||
slot_mapping = torch.arange(0, num_blocks * block_size, dtype=torch.int64)
|
||||
cpu_attn_reshape_and_cache(
|
||||
key=key_cache.view(-1, num_kv_heads, head_size),
|
||||
value=value_cache.view(-1, num_kv_heads, head_size),
|
||||
key_cache=packed_key_cache,
|
||||
value_cache=packed_value_cache,
|
||||
slot_mapping=slot_mapping,
|
||||
isa=isa,
|
||||
)
|
||||
|
||||
metadata = cpu_attn_get_scheduler_metadata(
|
||||
num_reqs=num_seqs,
|
||||
num_heads=num_query_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim=head_size,
|
||||
seq_lens=kv_lens_tensor,
|
||||
dtype=dtype,
|
||||
query_start_loc=cu_query_lens,
|
||||
causal=True,
|
||||
sliding_window_size=sliding_window if sliding_window is not None else -1,
|
||||
isa=isa,
|
||||
enable_kv_split=enable_kv_split,
|
||||
)
|
||||
|
||||
out_with_split = torch.empty_like(query)
|
||||
|
||||
def run_benchmark(iters: int) -> list[float]:
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
start_time = time.perf_counter_ns()
|
||||
cpu_attention_with_kv_cache(
|
||||
query=query,
|
||||
key_cache=packed_key_cache,
|
||||
value_cache=packed_value_cache,
|
||||
output=out_with_split,
|
||||
query_start_loc=cu_query_lens,
|
||||
seq_lens=kv_lens_tensor,
|
||||
scale=scale,
|
||||
causal=True,
|
||||
alibi_slopes=None,
|
||||
sliding_window=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=0,
|
||||
scheduler_metadata=metadata,
|
||||
s_aux=s_aux,
|
||||
)
|
||||
end_time = time.perf_counter_ns()
|
||||
times.append((end_time - start_time) / 1e6)
|
||||
return times
|
||||
|
||||
# warmup
|
||||
run_benchmark(5)
|
||||
# benchmark
|
||||
times = run_benchmark(iters)
|
||||
|
||||
time_min = min(times)
|
||||
time_max = max(times)
|
||||
time_mean = np.mean(times)
|
||||
time_std = np.std(times)
|
||||
|
||||
print("\tmin (ms) = ", time_min)
|
||||
print("\tmax (ms) = ", time_max)
|
||||
print("\tmean (ms) = ", time_mean)
|
||||
print("\tstd = ", time_std)
|
||||
print("\tmedian (ms) = ", np.median(times))
|
||||
|
||||
|
||||
def generate_seq_lens(
|
||||
batch_size: int,
|
||||
q_len_min: int,
|
||||
q_len_max: int,
|
||||
kv_len_min: int,
|
||||
kv_len_max: int,
|
||||
seed: int = 0,
|
||||
) -> list[tuple[int, int]]:
|
||||
assert 1 <= q_len_min <= q_len_max
|
||||
assert 1 <= kv_len_min <= kv_len_max
|
||||
assert kv_len_max >= q_len_min
|
||||
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
|
||||
def rint(lo: int, hi: int) -> int:
|
||||
return torch.randint(lo, hi + 1, (1,), generator=g).item()
|
||||
|
||||
seq_lens: list[tuple[int, int]] = []
|
||||
for _ in range(batch_size):
|
||||
# ensure q <= kv
|
||||
kv = rint(max(kv_len_min, q_len_min), kv_len_max)
|
||||
q = rint(q_len_min, min(q_len_max, kv))
|
||||
seq_lens.append((q, kv))
|
||||
|
||||
return seq_lens
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument("--q-len-min", type=int, default=512)
|
||||
parser.add_argument("--q-len-max", type=int, default=512)
|
||||
parser.add_argument("--kv-len-min", type=int, default=512)
|
||||
parser.add_argument("--kv-len-max", type=int, default=512)
|
||||
parser.add_argument("--num-blocks", type=int, default=4096)
|
||||
|
||||
parser.add_argument("--sliding-window", type=int, default=None)
|
||||
parser.add_argument("--num-query-heads", type=int, default=32)
|
||||
parser.add_argument("--num-kv-heads", type=int, default=8)
|
||||
parser.add_argument(
|
||||
"--head-size",
|
||||
type=int,
|
||||
choices=CPUAttentionBackend.get_supported_head_sizes(),
|
||||
default=128,
|
||||
)
|
||||
parser.add_argument("--enable-kv-split", action="store_true")
|
||||
parser.add_argument("--block-size", type=int, choices=[32, 64, 128], default=128)
|
||||
parser.add_argument(
|
||||
"--dtype", type=str, choices=["half", "bfloat16", "float"], default="bfloat16"
|
||||
)
|
||||
parser.add_argument("--use-sink", action="store_true")
|
||||
parser.add_argument(
|
||||
"--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--iters", type=int, default=20)
|
||||
|
||||
args = parser.parse_args()
|
||||
print(args)
|
||||
|
||||
seq_lens = generate_seq_lens(
|
||||
args.batch_size,
|
||||
args.q_len_min,
|
||||
args.q_len_max,
|
||||
args.kv_len_min,
|
||||
args.kv_len_max,
|
||||
args.seed,
|
||||
)
|
||||
|
||||
print("batch (query len, kv len) = ", seq_lens)
|
||||
|
||||
main(
|
||||
seq_lens=seq_lens,
|
||||
num_heads=(args.num_query_heads, args.num_kv_heads),
|
||||
head_size=args.head_size,
|
||||
sliding_window=args.sliding_window,
|
||||
dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
|
||||
block_size=args.block_size,
|
||||
num_blocks=args.num_blocks,
|
||||
use_sink=args.use_sink,
|
||||
enable_kv_split=args.enable_kv_split,
|
||||
isa=args.isa
|
||||
if args.isa is not None
|
||||
else get_attn_isa(args.block_size, STR_DTYPE_TO_TORCH_DTYPE[args.dtype]),
|
||||
seed=args.seed,
|
||||
iters=args.iters,
|
||||
)
|
||||
@@ -0,0 +1,175 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Check if CPU MoE operations are available
|
||||
try:
|
||||
from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
|
||||
except (ImportError, AttributeError) as e:
|
||||
print("ERROR: CPU fused MoE operations are not available on this platform.")
|
||||
print("This benchmark requires x86 CPU with proper vLLM CPU extensions compiled.")
|
||||
print(
|
||||
"The cpu_fused_moe kernel is typically available on Linux x86_64 "
|
||||
"with AVX2/AVX512."
|
||||
)
|
||||
print(f"Import error: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# ISA selection following test_cpu_fused_moe.py pattern
|
||||
ISA_CHOICES = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main(
|
||||
batch_size: int,
|
||||
expert_num: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
topk_num: int,
|
||||
use_bias: bool = False,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
activation: str = "silu",
|
||||
isa: str = "vec",
|
||||
seed: int = 0,
|
||||
iters: int = 20,
|
||||
) -> None:
|
||||
current_platform.seed_everything(seed)
|
||||
# up_dim = 2 * intermediate_size for gate + up projection
|
||||
up_dim = 2 * intermediate_size
|
||||
|
||||
input_tensor = torch.randn((batch_size, hidden_size), dtype=dtype) / (
|
||||
0.5 * hidden_size**0.5
|
||||
)
|
||||
|
||||
w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
|
||||
0.5 * hidden_size**0.5
|
||||
)
|
||||
w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
|
||||
0.5 * intermediate_size**0.5
|
||||
)
|
||||
|
||||
w13_bias = None
|
||||
w2_bias = None
|
||||
if use_bias:
|
||||
w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
|
||||
w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
|
||||
0.5 * hidden_size**0.5
|
||||
)
|
||||
|
||||
router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
|
||||
score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
|
||||
topk_weights, topk_ids = torch.topk(score, topk_num)
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
packed_w13 = cpu_prepack_moe_weight(w13, isa)
|
||||
packed_w2 = cpu_prepack_moe_weight(w2, isa)
|
||||
|
||||
def run_benchmark(iters: int) -> list[float]:
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
start_time = time.perf_counter_ns()
|
||||
_ = cpu_fused_moe(
|
||||
input_tensor,
|
||||
packed_w13,
|
||||
packed_w2,
|
||||
w13_bias,
|
||||
w2_bias,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation,
|
||||
isa,
|
||||
)
|
||||
end_time = time.perf_counter_ns()
|
||||
times.append((end_time - start_time) / 1e6)
|
||||
return times
|
||||
|
||||
# warmup
|
||||
run_benchmark(5)
|
||||
# benchmark
|
||||
times = run_benchmark(iters)
|
||||
|
||||
if not times:
|
||||
print("No iterations to measure. Set --iters > 0.")
|
||||
return
|
||||
|
||||
time_min = min(times)
|
||||
time_max = max(times)
|
||||
time_mean = np.mean(times)
|
||||
time_std = np.std(times)
|
||||
|
||||
print("\tmin (ms) = ", time_min)
|
||||
print("\tmax (ms) = ", time_max)
|
||||
print("\tmean (ms) = ", time_mean)
|
||||
print("\tstd = ", time_std)
|
||||
print("\tmedian (ms) = ", np.median(times))
|
||||
|
||||
# Calculate throughput metrics
|
||||
# FLOPs estimation: 2 * batch * topk * (hidden * up_dim + intermediate * hidden)
|
||||
flops_per_token = (
|
||||
2 * topk_num * (hidden_size * up_dim + intermediate_size * hidden_size)
|
||||
)
|
||||
total_flops = batch_size * flops_per_token
|
||||
tflops = total_flops / (time_mean * 1e-3) / 1e12
|
||||
print(f"\tthroughput (TFLOP/s) = {tflops:.4f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark the CPU fused MoE kernel.")
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument("--expert-num", type=int, default=8)
|
||||
parser.add_argument("--hidden-size", type=int, default=2880)
|
||||
parser.add_argument("--intermediate-size", type=int, default=2880)
|
||||
parser.add_argument(
|
||||
"--topk-num",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of experts to route each token to (default: expert_num // 2)",
|
||||
)
|
||||
parser.add_argument("--use-bias", action="store_true")
|
||||
parser.add_argument(
|
||||
"--activation",
|
||||
type=str,
|
||||
choices=["silu", "swigluoai"],
|
||||
default="silu",
|
||||
help="Activation function",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--isa",
|
||||
type=str,
|
||||
choices=ISA_CHOICES,
|
||||
default=ISA_CHOICES[0],
|
||||
help=f"ISA to use (available: {ISA_CHOICES})",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--iters", type=int, default=20)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Default topk_num to expert_num // 2, minimum 1
|
||||
topk_num = (
|
||||
args.topk_num if args.topk_num is not None else max(args.expert_num // 2, 1)
|
||||
)
|
||||
|
||||
print(args)
|
||||
|
||||
main(
|
||||
batch_size=args.batch_size,
|
||||
expert_num=args.expert_num,
|
||||
hidden_size=args.hidden_size,
|
||||
intermediate_size=args.intermediate_size,
|
||||
topk_num=topk_num,
|
||||
use_bias=args.use_bias,
|
||||
dtype=torch.bfloat16, # Following test_cpu_fused_moe.py
|
||||
activation=args.activation,
|
||||
isa=args.isa,
|
||||
seed=args.seed,
|
||||
iters=args.iters,
|
||||
)
|
||||
@@ -31,10 +31,15 @@ if(NOT qutlass_SOURCE_DIR)
|
||||
endif()
|
||||
message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
|
||||
|
||||
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0a" "${CUDA_ARCHS}")
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND QUTLASS_ARCHS)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
if(QUTLASS_ARCHS MATCHES "10\\.0a")
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
|
||||
if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)")
|
||||
set(QUTLASS_TARGET_CC 100)
|
||||
elseif(QUTLASS_ARCHS MATCHES "12\\.0a")
|
||||
set(QUTLASS_TARGET_CC 120)
|
||||
|
||||
@@ -38,7 +38,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG 86f8f157cf82aa2342743752b97788922dd7de43
|
||||
GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -27,6 +27,13 @@ void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
|
||||
const std::string& kv_cache_dtype,
|
||||
torch::Tensor& scale);
|
||||
|
||||
// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
|
||||
void concat_and_cache_mla_rope_fused(
|
||||
torch::Tensor& positions, torch::Tensor& q_pe, torch::Tensor& k_pe,
|
||||
torch::Tensor& kv_c, torch::Tensor& rope_cos_sin_cache, bool rope_is_neox,
|
||||
torch::Tensor& kv_cache_slot_mapping, torch::Tensor& kv_cache,
|
||||
const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale);
|
||||
|
||||
// Just for unittest
|
||||
void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
|
||||
const double scale, const std::string& kv_cache_dtype);
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "cuda_compat.h"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
#include "quantization/w8a8/fp8/common.cuh"
|
||||
#ifdef USE_ROCM
|
||||
#include "quantization/w8a8/fp8/amd/quant_utils.cuh"
|
||||
#else
|
||||
#include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
|
||||
#endif
|
||||
|
||||
#ifdef USE_ROCM
|
||||
#include <hip/hip_bf16.h>
|
||||
typedef __hip_bfloat16 __nv_bfloat16;
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's
|
||||
// using u16 as the backing type.
|
||||
template <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
|
||||
typename cache_t, Fp8KVCacheDataType kv_dt>
|
||||
__global__ void concat_and_cache_mla_rope_fused_kernel(
|
||||
const int64_t* __restrict__ positions, // [num_tokens]
|
||||
qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim]
|
||||
qk_t* __restrict__ k_pe, // [num_tokens, rot_dim]
|
||||
const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
|
||||
const qk_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
|
||||
// rot_dim // 2]
|
||||
const int rot_dim, const int64_t q_pe_stride_token,
|
||||
const int64_t q_pe_stride_head, const int64_t k_pe_stride,
|
||||
const int64_t kv_c_stride, const int num_q_heads,
|
||||
cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank +
|
||||
// rot_dim)]
|
||||
const int64_t* __restrict__ kv_cache_slot_mapping, // [num_tokens]
|
||||
const int block_stride, const int entry_stride, const int kv_lora_rank,
|
||||
const int block_size, const float* kv_cache_quant_scale) {
|
||||
// Each thread block is responsible for one token.
|
||||
const int64_t token_idx = blockIdx.x;
|
||||
const int64_t pos = positions[token_idx];
|
||||
|
||||
const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
|
||||
|
||||
const int embed_dim = rot_dim / 2;
|
||||
|
||||
// Q ROPE
|
||||
const int nq = num_q_heads * embed_dim;
|
||||
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
|
||||
int head_idx = i / embed_dim;
|
||||
int pair_idx = i % embed_dim;
|
||||
|
||||
// NOTE: Would be nice to have interleaved sin/cos so we could just load
|
||||
// both at the same time.
|
||||
qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
|
||||
qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
|
||||
|
||||
qk_t* q_pe_head_ptr =
|
||||
q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head;
|
||||
|
||||
int pair_idx_x, pair_idx_y;
|
||||
if constexpr (IS_NEOX) {
|
||||
// GPT-NeoX style rotary embedding.
|
||||
pair_idx_x = pair_idx;
|
||||
pair_idx_y = embed_dim + pair_idx;
|
||||
} else {
|
||||
// GPT-J style rotary embedding.
|
||||
pair_idx_x = pair_idx * 2;
|
||||
pair_idx_y = pair_idx * 2 + 1;
|
||||
}
|
||||
|
||||
qk_t x_src = q_pe_head_ptr[pair_idx_x];
|
||||
qk_t y_src = q_pe_head_ptr[pair_idx_y];
|
||||
|
||||
qk_t x_dst = x_src * cos - y_src * sin;
|
||||
qk_t y_dst = y_src * cos + x_src * sin;
|
||||
|
||||
q_pe_head_ptr[pair_idx_x] = x_dst;
|
||||
q_pe_head_ptr[pair_idx_y] = y_dst;
|
||||
}
|
||||
|
||||
const int64_t slot_idx = kv_cache_slot_mapping[token_idx];
|
||||
const int64_t block_idx = slot_idx / block_size;
|
||||
const int64_t entry_idx = slot_idx % block_size;
|
||||
|
||||
// NOTE: slot_idx can be -1 if the token is padded
|
||||
if (slot_idx < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// K with 1 HEAD
|
||||
for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) {
|
||||
int pair_idx = i;
|
||||
|
||||
qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
|
||||
qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
|
||||
|
||||
qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;
|
||||
|
||||
int pair_idx_x, pair_idx_y;
|
||||
if constexpr (IS_NEOX) {
|
||||
// GPT-NeoX style rotary embedding.
|
||||
pair_idx_x = pair_idx;
|
||||
pair_idx_y = embed_dim + pair_idx;
|
||||
} else {
|
||||
// GPT-J style rotary embedding.
|
||||
pair_idx_x = pair_idx * 2;
|
||||
pair_idx_y = pair_idx * 2 + 1;
|
||||
}
|
||||
|
||||
qk_t x_src = k_pe_head_ptr[pair_idx_x];
|
||||
qk_t y_src = k_pe_head_ptr[pair_idx_y];
|
||||
|
||||
qk_t x_dst = x_src * cos - y_src * sin;
|
||||
qk_t y_dst = y_src * cos + x_src * sin;
|
||||
|
||||
k_pe_head_ptr[pair_idx_x] = x_dst;
|
||||
k_pe_head_ptr[pair_idx_y] = y_dst;
|
||||
|
||||
// NOTE Why is this monster necessary?
|
||||
// When K is of type float16, the actual template replacement for
|
||||
// raw_kv_scalar_t with be u16. That's why it's used at the last moment
|
||||
// otherwise CUDA ALU would break.
|
||||
const raw_kv_scalar_t raw_x_value =
|
||||
*reinterpret_cast<const raw_kv_scalar_t*>(&x_dst);
|
||||
const raw_kv_scalar_t raw_y_value =
|
||||
*reinterpret_cast<const raw_kv_scalar_t*>(&y_dst);
|
||||
|
||||
cache_t* kv_cache_ptr = kv_cache + block_idx * block_stride +
|
||||
entry_idx * entry_stride + kv_lora_rank;
|
||||
|
||||
// MLA Cache Store
|
||||
if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
|
||||
kv_cache_ptr[pair_idx_x] = raw_x_value;
|
||||
kv_cache_ptr[pair_idx_y] = raw_y_value;
|
||||
} else {
|
||||
kv_cache_ptr[pair_idx_x] =
|
||||
fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>(
|
||||
raw_x_value, *kv_cache_quant_scale);
|
||||
kv_cache_ptr[pair_idx_y] =
|
||||
fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>(
|
||||
raw_y_value, *kv_cache_quant_scale);
|
||||
}
|
||||
}
|
||||
|
||||
// NOPE
|
||||
for (int i = threadIdx.x; i < kv_lora_rank; i += blockDim.x) {
|
||||
const qk_t* src_ptr = kv_c + token_idx * kv_c_stride + i;
|
||||
const raw_kv_scalar_t src_value =
|
||||
*reinterpret_cast<const raw_kv_scalar_t*>(src_ptr);
|
||||
|
||||
cache_t* kv_cache_ptr =
|
||||
kv_cache + block_idx * block_stride + entry_idx * entry_stride;
|
||||
|
||||
if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
|
||||
kv_cache_ptr[i] = src_value;
|
||||
} else {
|
||||
kv_cache_ptr[i] = fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>(
|
||||
src_value, *kv_cache_quant_scale);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
|
||||
do { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
|
||||
using qk_t = scalar_t; \
|
||||
if (rope_is_neox) { \
|
||||
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, true, RAW_KV_T, \
|
||||
CACHE_T, KV_DTYPE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
|
||||
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
|
||||
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
|
||||
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
|
||||
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
|
||||
entry_stride, kv_lora_rank, block_size, \
|
||||
kv_cache_quant_scale.data_ptr<float>()); \
|
||||
} else { \
|
||||
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T, \
|
||||
CACHE_T, KV_DTYPE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
|
||||
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
|
||||
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
|
||||
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
|
||||
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
|
||||
entry_stride, kv_lora_rank, block_size, \
|
||||
kv_cache_quant_scale.data_ptr<float>()); \
|
||||
} \
|
||||
}); \
|
||||
} while (false)
|
||||
|
||||
// Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
|
||||
// q_pe and k_pe are modified in place.
|
||||
// Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and
|
||||
// concat_and_cache_mla.
|
||||
void concat_and_cache_mla_rope_fused(
|
||||
torch::Tensor& positions, // [num_tokens]
|
||||
torch::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim]
|
||||
torch::Tensor& k_pe, // [num_tokens, rot_dim]
|
||||
torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
|
||||
torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
|
||||
bool rope_is_neox,
|
||||
torch::Tensor&
|
||||
kv_cache_slot_mapping, // [num_tokens] or [num_actual_tokens]
|
||||
torch::Tensor&
|
||||
kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
|
||||
const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
|
||||
const int64_t num_tokens = q_pe.size(0);
|
||||
|
||||
const int num_q_heads = q_pe.size(1);
|
||||
const int rot_dim = q_pe.size(2);
|
||||
const int kv_lora_rank = kv_c.size(1);
|
||||
|
||||
TORCH_CHECK(positions.size(0) >=
|
||||
num_tokens); // CUDA Graphs might pad this for us
|
||||
TORCH_CHECK_EQ(positions.dim(), 1);
|
||||
TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
|
||||
|
||||
TORCH_CHECK_EQ(q_pe.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
|
||||
TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
|
||||
TORCH_CHECK_EQ(q_pe.dim(), 3);
|
||||
|
||||
TORCH_CHECK_EQ(k_pe.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
|
||||
TORCH_CHECK_EQ(k_pe.dim(), 2);
|
||||
TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
|
||||
|
||||
TORCH_CHECK_EQ(kv_c.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
|
||||
TORCH_CHECK_EQ(kv_c.dim(), 2);
|
||||
TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
|
||||
TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
|
||||
|
||||
TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
|
||||
TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type());
|
||||
|
||||
TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long);
|
||||
|
||||
TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
|
||||
TORCH_CHECK_EQ(kv_cache.dim(), 3);
|
||||
|
||||
TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1);
|
||||
TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float);
|
||||
|
||||
int64_t q_pe_stride_token = q_pe.stride(0);
|
||||
int64_t q_pe_stride_head = q_pe.stride(1);
|
||||
|
||||
int64_t k_pe_stride = k_pe.stride(0);
|
||||
int64_t kv_c_stride = kv_c.stride(0);
|
||||
|
||||
int block_size = kv_cache.size(1);
|
||||
|
||||
int block_stride = kv_cache.stride(0);
|
||||
int entry_stride = kv_cache.stride(1);
|
||||
|
||||
int rope_block_size = std::min(num_q_heads * rot_dim / 2, 512);
|
||||
int mla_block_size = kv_lora_rank;
|
||||
int thread_block_size =
|
||||
std::min(std::max(rope_block_size, mla_block_size), 512);
|
||||
|
||||
dim3 grid(num_tokens, 1, 1);
|
||||
dim3 block(thread_block_size, 1, 1);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(positions));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
|
||||
CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED);
|
||||
}
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
#ifdef __aarch64__
|
||||
#include "cpu_attn_neon.hpp"
|
||||
// NEON requires head_dim to be a multiple of 32
|
||||
#define NEON_DISPATCH(...) \
|
||||
case cpu_attention::ISA::NEON: { \
|
||||
using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::NEON, \
|
||||
@@ -36,7 +37,9 @@
|
||||
switch (HEAD_DIM) { \
|
||||
CPU_ATTN_DISPATCH_CASE(32, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(64, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(80, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(96, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(112, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(128, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(160, __VA_ARGS__) \
|
||||
CPU_ATTN_DISPATCH_CASE(192, __VA_ARGS__) \
|
||||
|
||||
@@ -377,7 +377,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
const int32_t q_heads_per_kv, const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, const float scale) {
|
||||
constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t);
|
||||
static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
|
||||
// static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
|
||||
constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES;
|
||||
constexpr int64_t head_elem_num_pre_block =
|
||||
AMX_TILE_ROW_BYTES / sizeof(scalar_t);
|
||||
|
||||
@@ -264,7 +264,7 @@ class AttentionImpl<ISA::NEON, scalar_t, 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);
|
||||
// static_assert(HeadDim % HeadDimAlignment == 0);
|
||||
// the gemm micro kernel is Mx8
|
||||
static_assert(HeadDimAlignment % 8 == 0);
|
||||
static_assert(BlockSizeAlignment % 8 == 0);
|
||||
|
||||
@@ -457,8 +457,8 @@ __device__ inline T apply_scoring(T val) {
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, ScoringFunc SF>
|
||||
__device__ void topk_with_k2(T* output, T const* input, T const* bias,
|
||||
template <typename T, typename BiasT, ScoringFunc SF>
|
||||
__device__ void topk_with_k2(T* output, T const* input, BiasT const* bias,
|
||||
cg::thread_block_tile<32> const& tile,
|
||||
int32_t const lane_id,
|
||||
int const num_experts_per_group) {
|
||||
@@ -469,7 +469,7 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
|
||||
if (num_experts_per_group > WARP_SIZE) {
|
||||
for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
|
||||
T value = apply_scoring<SF>(input[i]);
|
||||
value = value + bias[i];
|
||||
value = value + static_cast<T>(bias[i]);
|
||||
|
||||
if (value > largest) {
|
||||
second_largest = largest;
|
||||
@@ -481,7 +481,7 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
|
||||
} else {
|
||||
for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
|
||||
T value = apply_scoring<SF>(input[i]);
|
||||
value = value + bias[i];
|
||||
value = value + static_cast<T>(bias[i]);
|
||||
largest = value;
|
||||
}
|
||||
}
|
||||
@@ -503,8 +503,8 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, ScoringFunc SF>
|
||||
__global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
|
||||
template <typename T, typename BiasT, ScoringFunc SF>
|
||||
__global__ void topk_with_k2_kernel(T* output, T* input, BiasT const* bias,
|
||||
int64_t const num_tokens,
|
||||
int64_t const num_cases,
|
||||
int64_t const n_group,
|
||||
@@ -517,7 +517,7 @@ __global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
|
||||
input += case_id * num_experts_per_group;
|
||||
// bias is per expert group, offset to current group
|
||||
int32_t group_id = case_id % n_group;
|
||||
T const* group_bias = bias + group_id * num_experts_per_group;
|
||||
BiasT const* group_bias = bias + group_id * num_experts_per_group;
|
||||
output += case_id;
|
||||
|
||||
cg::thread_block block = cg::this_thread_block();
|
||||
@@ -526,18 +526,19 @@ __global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
topk_with_k2<T, SF>(output, input, group_bias, tile, lane_id,
|
||||
num_experts_per_group);
|
||||
topk_with_k2<T, BiasT, SF>(output, input, group_bias, tile, lane_id,
|
||||
num_experts_per_group);
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T, typename IdxT, ScoringFunc SF, int NGroup = -1>
|
||||
template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
|
||||
int NGroup = -1>
|
||||
__global__ void group_idx_and_topk_idx_kernel(
|
||||
T* scores, T const* group_scores, float* topk_values, IdxT* topk_indices,
|
||||
T const* bias, int64_t const num_tokens, int64_t const n_group,
|
||||
BiasT const* bias, int64_t const num_tokens, int64_t const n_group,
|
||||
int64_t const topk_group, int64_t const topk, int64_t const num_experts,
|
||||
int64_t const num_experts_per_group, bool renormalize,
|
||||
double routed_scaling_factor) {
|
||||
@@ -623,7 +624,7 @@ __global__ void group_idx_and_topk_idx_kernel(
|
||||
T input = scores[offset + i];
|
||||
if (is_finite(input)) {
|
||||
T score = apply_scoring<SF>(input);
|
||||
candidates = score + bias[offset + i];
|
||||
candidates = score + static_cast<T>(bias[offset + i]);
|
||||
}
|
||||
}
|
||||
queue.add(candidates, offset + i);
|
||||
@@ -698,10 +699,10 @@ __global__ void group_idx_and_topk_idx_kernel(
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T, typename IdxT, ScoringFunc SF>
|
||||
template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
|
||||
inline void launch_group_idx_and_topk_kernel(
|
||||
cudaLaunchConfig_t const& config, T* scores, T* group_scores,
|
||||
float* topk_values, IdxT* topk_indices, T const* bias,
|
||||
float* topk_values, IdxT* topk_indices, BiasT const* bias,
|
||||
int64_t const num_tokens, int64_t const n_group, int64_t const topk_group,
|
||||
int64_t const topk, int64_t const num_experts,
|
||||
int64_t const num_experts_per_group, bool const renormalize,
|
||||
@@ -715,36 +716,36 @@ inline void launch_group_idx_and_topk_kernel(
|
||||
|
||||
switch (n_group) {
|
||||
case 4: {
|
||||
launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 4>);
|
||||
launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 4>);
|
||||
break;
|
||||
}
|
||||
case 8: {
|
||||
launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 8>);
|
||||
launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 8>);
|
||||
break;
|
||||
}
|
||||
case 16: {
|
||||
launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 16>);
|
||||
launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 16>);
|
||||
break;
|
||||
}
|
||||
case 32: {
|
||||
launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 32>);
|
||||
launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 32>);
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF>);
|
||||
launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF>);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IdxT>
|
||||
template <typename T, typename BiasT, typename IdxT>
|
||||
void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
|
||||
IdxT* topk_indices, T const* bias, int64_t const num_tokens,
|
||||
int64_t const num_experts, int64_t const n_group,
|
||||
int64_t const topk_group, int64_t const topk,
|
||||
bool const renormalize, double const routed_scaling_factor,
|
||||
int const scoring_func, bool enable_pdl = false,
|
||||
cudaStream_t const stream = 0) {
|
||||
IdxT* topk_indices, BiasT const* bias,
|
||||
int64_t const num_tokens, int64_t const num_experts,
|
||||
int64_t const n_group, int64_t const topk_group,
|
||||
int64_t const topk, bool const renormalize,
|
||||
double const routed_scaling_factor, int const scoring_func,
|
||||
bool enable_pdl = false, cudaStream_t const stream = 0) {
|
||||
int64_t num_cases = num_tokens * n_group;
|
||||
int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;
|
||||
cudaLaunchConfig_t config;
|
||||
@@ -765,12 +766,12 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
|
||||
};
|
||||
switch (sf) {
|
||||
case SCORING_NONE: {
|
||||
auto* kernel_instance1 = &topk_with_k2_kernel<T, SCORING_NONE>;
|
||||
auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_NONE>;
|
||||
launch_topk_with_k2(kernel_instance1);
|
||||
break;
|
||||
}
|
||||
case SCORING_SIGMOID: {
|
||||
auto* kernel_instance1 = &topk_with_k2_kernel<T, SCORING_SIGMOID>;
|
||||
auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_SIGMOID>;
|
||||
launch_topk_with_k2(kernel_instance1);
|
||||
break;
|
||||
}
|
||||
@@ -794,14 +795,14 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
|
||||
config.attrs = attrs;
|
||||
switch (sf) {
|
||||
case SCORING_NONE: {
|
||||
launch_group_idx_and_topk_kernel<T, IdxT, SCORING_NONE>(
|
||||
launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_NONE>(
|
||||
config, scores, group_scores, topk_values, topk_indices, bias,
|
||||
num_tokens, n_group, topk_group, topk, num_experts,
|
||||
num_experts_per_group, renormalize, routed_scaling_factor);
|
||||
break;
|
||||
}
|
||||
case SCORING_SIGMOID: {
|
||||
launch_group_idx_and_topk_kernel<T, IdxT, SCORING_SIGMOID>(
|
||||
launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_SIGMOID>(
|
||||
config, scores, group_scores, topk_values, topk_indices, bias,
|
||||
num_tokens, n_group, topk_group, topk, num_experts,
|
||||
num_experts_per_group, renormalize, routed_scaling_factor);
|
||||
@@ -812,17 +813,23 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
|
||||
}
|
||||
}
|
||||
|
||||
#define INSTANTIATE_NOAUX_TC(T, IdxT) \
|
||||
template void invokeNoAuxTc<T, IdxT>( \
|
||||
T * scores, T * group_scores, float* topk_values, IdxT* topk_indices, \
|
||||
T const* bias, int64_t const num_tokens, int64_t const num_experts, \
|
||||
int64_t const n_group, int64_t const topk_group, int64_t const topk, \
|
||||
bool const renormalize, double const routed_scaling_factor, \
|
||||
#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT) \
|
||||
template void invokeNoAuxTc<T, BiasT, IdxT>( \
|
||||
T * scores, T * group_scores, float* topk_values, IdxT* topk_indices, \
|
||||
BiasT const* bias, int64_t const num_tokens, int64_t const num_experts, \
|
||||
int64_t const n_group, int64_t const topk_group, int64_t const topk, \
|
||||
bool const renormalize, double const routed_scaling_factor, \
|
||||
int const scoring_func, bool enable_pdl, cudaStream_t const stream);
|
||||
|
||||
INSTANTIATE_NOAUX_TC(float, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(half, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(__nv_bfloat16, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(float, float, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(float, half, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(half, float, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(half, half, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t);
|
||||
INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t);
|
||||
} // end namespace moe
|
||||
} // namespace vllm
|
||||
|
||||
@@ -831,6 +838,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
|
||||
int64_t topk, bool renormalize, double routed_scaling_factor,
|
||||
torch::Tensor const& bias, int64_t scoring_func = 0) {
|
||||
auto data_type = scores.scalar_type();
|
||||
auto bias_type = bias.scalar_type();
|
||||
auto input_size = scores.sizes();
|
||||
int64_t num_tokens = input_size[0];
|
||||
int64_t num_experts = input_size[1];
|
||||
@@ -854,39 +862,62 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
|
||||
|
||||
auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
|
||||
|
||||
#define LAUNCH_KERNEL(T, IdxT) \
|
||||
do { \
|
||||
switch (bias_type) { \
|
||||
case torch::kFloat16: \
|
||||
vllm::moe::invokeNoAuxTc<T, half, IdxT>( \
|
||||
reinterpret_cast<T*>(scores.mutable_data_ptr()), \
|
||||
reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \
|
||||
reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
|
||||
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
|
||||
reinterpret_cast<half const*>(bias.data_ptr()), num_tokens, \
|
||||
num_experts, n_group, topk_group, topk, renormalize, \
|
||||
routed_scaling_factor, static_cast<int>(scoring_func), false, \
|
||||
stream); \
|
||||
break; \
|
||||
case torch::kFloat32: \
|
||||
vllm::moe::invokeNoAuxTc<T, float, IdxT>( \
|
||||
reinterpret_cast<T*>(scores.mutable_data_ptr()), \
|
||||
reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \
|
||||
reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
|
||||
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
|
||||
reinterpret_cast<float const*>(bias.data_ptr()), num_tokens, \
|
||||
num_experts, n_group, topk_group, topk, renormalize, \
|
||||
routed_scaling_factor, static_cast<int>(scoring_func), false, \
|
||||
stream); \
|
||||
break; \
|
||||
case torch::kBFloat16: \
|
||||
vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>( \
|
||||
reinterpret_cast<T*>(scores.mutable_data_ptr()), \
|
||||
reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \
|
||||
reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
|
||||
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
|
||||
reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), \
|
||||
num_tokens, num_experts, n_group, topk_group, topk, renormalize, \
|
||||
routed_scaling_factor, static_cast<int>(scoring_func), false, \
|
||||
stream); \
|
||||
break; \
|
||||
default: \
|
||||
throw std::invalid_argument( \
|
||||
"Invalid bias dtype, only supports float16, float32, and " \
|
||||
"bfloat16"); \
|
||||
break; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
switch (data_type) {
|
||||
case torch::kFloat16:
|
||||
// Handle Float16
|
||||
vllm::moe::invokeNoAuxTc<half, int32_t>(
|
||||
reinterpret_cast<half*>(scores.mutable_data_ptr()),
|
||||
reinterpret_cast<half*>(group_scores.mutable_data_ptr()),
|
||||
reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
|
||||
reinterpret_cast<half const*>(bias.data_ptr()), num_tokens,
|
||||
num_experts, n_group, topk_group, topk, renormalize,
|
||||
routed_scaling_factor, static_cast<int>(scoring_func), false, stream);
|
||||
LAUNCH_KERNEL(half, int32_t);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
// Handle Float32
|
||||
vllm::moe::invokeNoAuxTc<float, int32_t>(
|
||||
reinterpret_cast<float*>(scores.mutable_data_ptr()),
|
||||
reinterpret_cast<float*>(group_scores.mutable_data_ptr()),
|
||||
reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
|
||||
reinterpret_cast<float const*>(bias.data_ptr()), num_tokens,
|
||||
num_experts, n_group, topk_group, topk, renormalize,
|
||||
routed_scaling_factor, static_cast<int>(scoring_func), false, stream);
|
||||
LAUNCH_KERNEL(float, int32_t);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
// Handle BFloat16
|
||||
vllm::moe::invokeNoAuxTc<__nv_bfloat16, int32_t>(
|
||||
reinterpret_cast<__nv_bfloat16*>(scores.mutable_data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16*>(group_scores.mutable_data_ptr()),
|
||||
reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), num_tokens,
|
||||
num_experts, n_group, topk_group, topk, renormalize,
|
||||
routed_scaling_factor, static_cast<int>(scoring_func), false, stream);
|
||||
LAUNCH_KERNEL(__nv_bfloat16, int32_t);
|
||||
break;
|
||||
default:
|
||||
// Handle other data types
|
||||
@@ -894,5 +925,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
|
||||
"Invalid dtype, only supports float16, float32, and bfloat16");
|
||||
break;
|
||||
}
|
||||
#undef LAUNCH_KERNEL
|
||||
return {topk_values, topk_indices};
|
||||
}
|
||||
|
||||
@@ -7,20 +7,20 @@
|
||||
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
|
||||
#include "core/scalar_type.hpp"
|
||||
|
||||
#define MARLIN_KERNEL_PARAMS \
|
||||
const int4 *__restrict__ A, const int4 *__restrict__ B, \
|
||||
int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
|
||||
const int4 *__restrict__ b_bias_ptr, \
|
||||
const float *__restrict__ a_scales_ptr, \
|
||||
const int4 *__restrict__ scales_ptr, \
|
||||
const uint16_t *__restrict__ global_scale_ptr, \
|
||||
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
|
||||
const int32_t *__restrict__ sorted_token_ids_ptr, \
|
||||
const int32_t *__restrict__ expert_ids_ptr, \
|
||||
const int32_t *__restrict__ num_tokens_past_padded_ptr, \
|
||||
const float *__restrict__ topk_weights_ptr, int top_k, \
|
||||
bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
|
||||
int prob_n, int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
|
||||
#define MARLIN_KERNEL_PARAMS \
|
||||
const int4 *__restrict__ A, const int4 *__restrict__ B, \
|
||||
int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
|
||||
const int4 *__restrict__ b_bias_ptr, \
|
||||
const float *__restrict__ a_scales_ptr, \
|
||||
const int4 *__restrict__ scales_ptr, \
|
||||
const uint16_t *__restrict__ global_scale_ptr, \
|
||||
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
|
||||
const int32_t *__restrict__ sorted_token_ids_ptr, \
|
||||
const int32_t *__restrict__ expert_ids_ptr, \
|
||||
const int32_t *__restrict__ num_tokens_past_padded_ptr, \
|
||||
const float *__restrict__ topk_weights_ptr, int top_k, \
|
||||
bool mul_topk_weights, int num_groups, int prob_m, int prob_n, \
|
||||
int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
|
||||
bool use_fp32_reduce
|
||||
|
||||
namespace MARLIN_NAMESPACE_NAME {
|
||||
|
||||
@@ -71,7 +71,6 @@ __global__ void Marlin(
|
||||
const float* __restrict__ topk_weights_ptr, // moe top weights
|
||||
int top_k, // num of experts per token
|
||||
bool mul_topk_weights, // mul topk weights or not
|
||||
bool is_ep, // expert parallelism
|
||||
int num_groups, // number of scale groups per output channel
|
||||
int prob_m, // batch dimension m
|
||||
int prob_n, // output dimension n
|
||||
@@ -273,7 +272,6 @@ __global__ void Marlin(
|
||||
const float* __restrict__ topk_weights_ptr, // moe top weights
|
||||
int top_k, // num of experts per token
|
||||
bool mul_topk_weights, // mul topk weights or not
|
||||
bool is_ep, // expert parallelism
|
||||
int num_groups, // number of scale groups per output channel
|
||||
int prob_m, // batch dimension m
|
||||
int prob_n, // output dimension n
|
||||
@@ -376,14 +374,6 @@ __global__ void Marlin(
|
||||
|
||||
// parallel: num valid moe blocks
|
||||
int parallel = num_tokens_past_padded / moe_block_size;
|
||||
int num_valid_blocks = parallel;
|
||||
if (is_ep) {
|
||||
for (int i = 0; i < parallel; i++) {
|
||||
if (expert_ids_ptr[i] == -1) num_valid_blocks--;
|
||||
}
|
||||
}
|
||||
int num_invalid_blocks = parallel - num_valid_blocks;
|
||||
parallel = num_valid_blocks;
|
||||
|
||||
int k_tiles = prob_k / 16 / thread_k_blocks;
|
||||
int n_tiles = prob_n / 16 / thread_n_blocks;
|
||||
@@ -538,22 +528,8 @@ __global__ void Marlin(
|
||||
if (par_id >= parallel) return;
|
||||
|
||||
old_expert_id = expert_id;
|
||||
if (num_invalid_blocks > 0) {
|
||||
int skip_count = par_id;
|
||||
for (int i = 0; i < num_tokens_past_padded / moe_block_size; i++) {
|
||||
expert_id = expert_ids_ptr[i];
|
||||
if (expert_id != -1) {
|
||||
if (skip_count == 0) {
|
||||
block_id = i;
|
||||
break;
|
||||
};
|
||||
skip_count--;
|
||||
};
|
||||
}
|
||||
} else {
|
||||
block_id = par_id;
|
||||
expert_id = expert_ids_ptr[block_id];
|
||||
}
|
||||
block_id = par_id;
|
||||
expert_id = expert_ids_ptr[block_id];
|
||||
|
||||
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
|
||||
uint16_t val = global_scale_ptr[expert_id];
|
||||
|
||||
@@ -336,14 +336,14 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
|
||||
void* perm, void* a_tmp, void* sorted_token_ids,
|
||||
void* expert_ids, void* num_tokens_past_padded,
|
||||
void* topk_weights, int moe_block_size, int num_experts,
|
||||
int top_k, bool mul_topk_weights, bool is_ep, int prob_m,
|
||||
int prob_n, int prob_k, void* workspace,
|
||||
vllm::ScalarType const& a_type, vllm::ScalarType const& b_type,
|
||||
vllm::ScalarType const& c_type, vllm::ScalarType const& s_type,
|
||||
bool has_bias, bool has_act_order, bool is_k_full, bool has_zp,
|
||||
int num_groups, int group_size, int dev, cudaStream_t stream,
|
||||
int thread_k, int thread_n, int sms, int blocks_per_sm,
|
||||
bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
|
||||
int top_k, bool mul_topk_weights, int prob_m, int prob_n,
|
||||
int prob_k, void* workspace, vllm::ScalarType const& a_type,
|
||||
vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
|
||||
vllm::ScalarType const& s_type, bool has_bias,
|
||||
bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
|
||||
int group_size, int dev, cudaStream_t stream, int thread_k,
|
||||
int thread_n, int sms, int blocks_per_sm, bool use_atomic_add,
|
||||
bool use_fp32_reduce, bool is_zp_float) {
|
||||
int thread_m_blocks = div_ceil(moe_block_size, 16);
|
||||
bool m_block_size_8 = moe_block_size == 8;
|
||||
bool is_a_8bit = a_type.size_bits() == 8;
|
||||
@@ -523,7 +523,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
|
||||
kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
|
||||
A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr,
|
||||
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
|
||||
topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
|
||||
topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m,
|
||||
prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
|
||||
// clang-format on
|
||||
}
|
||||
@@ -541,7 +541,7 @@ torch::Tensor moe_wna16_marlin_gemm(
|
||||
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
|
||||
torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
|
||||
torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
|
||||
int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
|
||||
int64_t moe_block_size, int64_t top_k, bool mul_topk_weights,
|
||||
vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
|
||||
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
|
||||
bool is_zp_float, int64_t thread_k, int64_t thread_n,
|
||||
@@ -855,9 +855,9 @@ torch::Tensor moe_wna16_marlin_gemm(
|
||||
perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(),
|
||||
expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
|
||||
topk_weights.data_ptr(), moe_block_size, num_experts, top_k,
|
||||
mul_topk_weights, is_ep, size_m, size_n, size_k, workspace.data_ptr(),
|
||||
a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
|
||||
has_zp, num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
|
||||
mul_topk_weights, size_m, size_n, size_k, workspace.data_ptr(), a_type,
|
||||
b_type, c_type, s_type, has_bias, has_act_order, is_k_full, has_zp,
|
||||
num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
|
||||
thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
|
||||
is_zp_float);
|
||||
|
||||
@@ -866,4 +866,4 @@ torch::Tensor moe_wna16_marlin_gemm(
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -71,7 +71,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"Tensor sorted_token_ids,"
|
||||
"Tensor! expert_ids, Tensor! num_tokens_past_padded,"
|
||||
"Tensor! topk_weights, int moe_block_size, int top_k, "
|
||||
"bool mul_topk_weights, bool is_ep, int b_type_id,"
|
||||
"bool mul_topk_weights, int b_type_id,"
|
||||
"int size_m, int size_n, int size_k,"
|
||||
"bool is_full_k, bool use_atomic_add,"
|
||||
"bool use_fp32_reduce, bool is_zp_float,"
|
||||
|
||||
+15
-2
@@ -2,6 +2,7 @@
|
||||
|
||||
#include <optional>
|
||||
#include <torch/library.h>
|
||||
#include <tuple>
|
||||
|
||||
#include "core/scalar_type.hpp"
|
||||
|
||||
@@ -265,6 +266,11 @@ void get_cutlass_moe_mm_problem_sizes(
|
||||
const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
|
||||
std::optional<bool> force_swap_ab = std::nullopt);
|
||||
|
||||
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
|
||||
const torch::Tensor& expert_first_token_offset,
|
||||
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
|
||||
const int64_t n, const int64_t k, const bool swap_ab);
|
||||
|
||||
void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
|
||||
torch::Tensor& problem_sizes1,
|
||||
torch::Tensor& problem_sizes2,
|
||||
@@ -301,6 +307,12 @@ void scaled_fp4_experts_quant(
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts);
|
||||
|
||||
void silu_and_mul_scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts);
|
||||
|
||||
void per_token_group_quant_fp8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q, torch::Tensor& output_s,
|
||||
int64_t group_size, double eps, double fp8_min,
|
||||
@@ -335,8 +347,9 @@ torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
|
||||
|
||||
void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit);
|
||||
|
||||
void static_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
torch::Tensor const& scale);
|
||||
void static_scaled_fp8_quant(
|
||||
torch::Tensor& out, torch::Tensor const& input, torch::Tensor const& scale,
|
||||
std::optional<std::tuple<int64_t, int64_t>> group_shape = std::nullopt);
|
||||
|
||||
void dynamic_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
@@ -31,37 +31,6 @@
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// silu in float32
|
||||
__device__ __forceinline__ float silu(float x) {
|
||||
return __fdividef(x, (1.f + __expf(-x)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float2 silu2(float2 x) {
|
||||
return make_float2(silu(x.x), silu(x.y));
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__inline__ __device__ PackedVec<Type> compute_silu_mul(PackedVec<Type>& vec,
|
||||
PackedVec<Type>& vec2) {
|
||||
PackedVec<Type> result;
|
||||
using packed_type = typename TypeConverter<Type>::Type;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
|
||||
// silu_mul in float32
|
||||
if constexpr (std::is_same_v<Type, half>) {
|
||||
float2 silu_vec = silu2(__half22float2(vec.elts[i]));
|
||||
result.elts[i] =
|
||||
__float22half2_rn(__fmul2_rn(silu_vec, __half22float2(vec2.elts[i])));
|
||||
} else {
|
||||
float2 silu_vec = silu2(__bfloat1622float2(vec.elts[i]));
|
||||
result.elts[i] = __float22bfloat162_rn(
|
||||
__fmul2_rn(silu_vec, __bfloat1622float2(vec2.elts[i])));
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
|
||||
@@ -62,7 +62,9 @@ __global__ void __get_group_gemm_starts(
|
||||
ElementSF* a_scales_base_as_int, ElementSF* b_scales_base_as_int,
|
||||
ElementAccumulator* alphas_base_as_int, const int32_t* expert_offsets,
|
||||
const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
|
||||
const int K, const int N) {
|
||||
int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
|
||||
const int64_t a_stride_val, const int64_t b_stride_val,
|
||||
const int64_t c_stride_val, const int K, const int N) {
|
||||
int64_t expert_id = threadIdx.x;
|
||||
if (expert_id >= gridDim.x * blockDim.x) {
|
||||
return;
|
||||
@@ -103,6 +105,11 @@ __global__ void __get_group_gemm_starts(
|
||||
// Shape of alpha = [E]
|
||||
alpha_offsets[expert_id] = alphas_base_as_int + expert_id;
|
||||
|
||||
// Initialize strides (constant across all experts, avoids separate kernels)
|
||||
a_strides[expert_id] = a_stride_val;
|
||||
b_strides[expert_id] = b_stride_val;
|
||||
c_strides[expert_id] = c_stride_val;
|
||||
|
||||
LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
|
||||
LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
|
||||
|
||||
@@ -135,7 +142,11 @@ __global__ void __get_group_gemm_starts(
|
||||
static_cast<float*>(alphas.data_ptr()), \
|
||||
static_cast<int32_t*>(expert_offsets.data_ptr()), \
|
||||
static_cast<int32_t*>(sf_offsets.data_ptr()), \
|
||||
static_cast<int32_t*>(problem_sizes.data_ptr()), K, N); \
|
||||
static_cast<int32_t*>(problem_sizes.data_ptr()), \
|
||||
static_cast<int64_t*>(a_strides.data_ptr()), \
|
||||
static_cast<int64_t*>(b_strides.data_ptr()), \
|
||||
static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \
|
||||
b_stride_val, c_stride_val, K, N); \
|
||||
}
|
||||
|
||||
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
|
||||
@@ -144,6 +155,9 @@ void run_get_group_gemm_starts(
|
||||
const torch::Tensor& out_starts, const torch::Tensor& a_scales_starts,
|
||||
const torch::Tensor& b_scales_starts, const torch::Tensor& alpha_starts,
|
||||
const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
|
||||
const torch::Tensor& a_strides, const torch::Tensor& b_strides,
|
||||
const torch::Tensor& c_strides, int64_t a_stride_val, int64_t b_stride_val,
|
||||
int64_t c_stride_val,
|
||||
/*these are used for their base addresses*/
|
||||
torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
|
||||
torch::Tensor const& out_tensors, torch::Tensor const& a_scales,
|
||||
@@ -269,17 +283,16 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
|
||||
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
|
||||
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
|
||||
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
|
||||
torch::Tensor c_strides1 =
|
||||
torch::full({num_experts}, output.stride(0), options_int);
|
||||
torch::Tensor a_strides1 =
|
||||
torch::full({num_experts}, a.stride(0) * 2, options_int);
|
||||
torch::Tensor b_strides1 =
|
||||
torch::full({num_experts}, b.stride(1) * 2, options_int);
|
||||
torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
|
||||
torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
|
||||
torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
|
||||
|
||||
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
|
||||
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
|
||||
layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
|
||||
expert_offsets, sf_offsets, problem_sizes, M, N, K);
|
||||
layout_sfa, layout_sfb, a_strides1, b_strides1, c_strides1,
|
||||
a.stride(0) * 2, b.stride(1) * 2, output.stride(0), a, b, output,
|
||||
a_blockscale, b_blockscales, alphas, expert_offsets, sf_offsets,
|
||||
problem_sizes, M, N, K);
|
||||
|
||||
// Create an instance of the GEMM
|
||||
Gemm gemm_op;
|
||||
@@ -444,17 +457,16 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
|
||||
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
|
||||
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
|
||||
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
|
||||
torch::Tensor c_strides1 =
|
||||
torch::full({num_experts}, output.stride(0), options_int);
|
||||
torch::Tensor a_strides1 =
|
||||
torch::full({num_experts}, a.stride(0) * 2, options_int);
|
||||
torch::Tensor b_strides1 =
|
||||
torch::full({num_experts}, b.stride(1) * 2, options_int);
|
||||
torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
|
||||
torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
|
||||
torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
|
||||
|
||||
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
|
||||
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
|
||||
layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
|
||||
expert_offsets, sf_offsets, problem_sizes, M, N, K);
|
||||
layout_sfa, layout_sfb, a_strides1, b_strides1, c_strides1,
|
||||
a.stride(0) * 2, b.stride(1) * 2, output.stride(0), a, b, output,
|
||||
a_blockscale, b_blockscales, alphas, expert_offsets, sf_offsets,
|
||||
problem_sizes, M, N, K);
|
||||
|
||||
// Create an instance of the GEMM
|
||||
Gemm gemm_op;
|
||||
|
||||
@@ -31,8 +31,12 @@
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// NVFP4 quantization kernel for experts (low-latency path).
|
||||
// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
|
||||
// SiLU(gate)*up before quantization.
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
|
||||
template <class Type, bool FUSE_SILU_MUL = false, bool UE8M0_SF = false,
|
||||
bool SMALL_NUM_EXPERTS = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
float const* SFScale, uint32_t* out, uint32_t* SFout,
|
||||
@@ -50,6 +54,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
// When fusing SiLU+Mul, input has gate || up layout (doubled width)
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
// Each global thread processes one element
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
@@ -58,13 +64,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
int64_t inOffset = rowIdx * colsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
// Get the output tensor offset.
|
||||
// Same as inOffset because 8 elements are packed into one uint32_t.
|
||||
int64_t outOffset = inOffset;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
// Find index within the experts using different strategies based on expert
|
||||
// count
|
||||
int rowIdx_in_expert = 0;
|
||||
@@ -111,6 +110,23 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
}
|
||||
}
|
||||
|
||||
// Load input and optionally apply fused SiLU+Mul
|
||||
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
PackedVec quant_input;
|
||||
if constexpr (FUSE_SILU_MUL) {
|
||||
PackedVec in_vec_up =
|
||||
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
|
||||
quant_input = compute_silu_mul(in_vec, in_vec_up);
|
||||
} else {
|
||||
quant_input = in_vec;
|
||||
}
|
||||
|
||||
// Get the output tensor offset.
|
||||
// Same as inOffset because 8 elements are packed into one uint32_t.
|
||||
int64_t outOffset = rowIdx * colsPerRow + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
@@ -124,12 +140,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
// Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
|
||||
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
|
||||
// NVFP4 quantization kernel for LARGE_M_TOPK = true (large m_topk optimized
|
||||
// version). When FUSE_SILU_MUL=true, expects input with gate||up layout and
|
||||
// fuses SiLU(gate)*up before quantization.
|
||||
template <class Type, bool FUSE_SILU_MUL = false, bool UE8M0_SF = false,
|
||||
bool SMALL_NUM_EXPERTS = false>
|
||||
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
float const* SFScale, uint32_t* out, uint32_t* SFout,
|
||||
@@ -167,6 +187,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
// When fusing SiLU+Mul, input has gate || up layout (doubled width)
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
// Each global thread processes one element
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
@@ -175,11 +197,6 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
int64_t inOffset = rowIdx * colsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
int64_t outOffset = inOffset;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
// Find expert using binary search for better performance with large m_topk
|
||||
int rowIdx_in_expert = 0;
|
||||
int expert_idx = 0;
|
||||
@@ -204,6 +221,21 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
}
|
||||
}
|
||||
|
||||
// Load input and optionally apply fused SiLU+Mul
|
||||
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
PackedVec quant_input;
|
||||
if constexpr (FUSE_SILU_MUL) {
|
||||
PackedVec in_vec_up =
|
||||
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
|
||||
quant_input = compute_silu_mul(in_vec, in_vec_up);
|
||||
} else {
|
||||
quant_input = in_vec;
|
||||
}
|
||||
|
||||
int64_t outOffset = rowIdx * colsPerRow + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
|
||||
|
||||
uint32_t* SFout_in_expert =
|
||||
@@ -214,11 +246,12 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
template <typename T, bool FUSE_SILU_MUL = false>
|
||||
void quant_impl(void* output, void* output_scale, void* input,
|
||||
void* input_global_scale, void* input_offset_by_experts,
|
||||
void* output_scale_offset_by_experts, int m_topk, int k,
|
||||
@@ -246,7 +279,7 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
if (blockRepeat > 1) {
|
||||
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
|
||||
if (n_experts >= 4) {
|
||||
cvt_fp16_to_fp4<T, false, false>
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
@@ -256,34 +289,37 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts);
|
||||
} else {
|
||||
cvt_fp16_to_fp4<T, false, true><<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
reinterpret_cast<uint32_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts);
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, true>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
reinterpret_cast<uint32_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts);
|
||||
}
|
||||
} else {
|
||||
if (n_experts >= 16) {
|
||||
cvt_fp16_to_fp4<T, false, false><<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
reinterpret_cast<uint32_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts, /* bool low_latency */ true);
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
reinterpret_cast<uint32_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts, /* bool low_latency */ true);
|
||||
} else {
|
||||
cvt_fp16_to_fp4<T, false, true><<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
reinterpret_cast<uint32_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts, /* bool low_latency */ true);
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, true>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<float*>(input_global_scale),
|
||||
reinterpret_cast<uint32_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts, /* bool low_latency */ true);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -304,19 +340,19 @@ constexpr auto FLOAT = at::ScalarType::Float;
|
||||
constexpr auto INT = at::ScalarType::Int;
|
||||
constexpr auto UINT8 = at::ScalarType::Byte;
|
||||
|
||||
void scaled_fp4_experts_quant_sm1xxa(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
// Common validation for fp4 experts quantization entry points.
|
||||
static void validate_fp4_experts_quant_inputs(
|
||||
torch::Tensor const& output, torch::Tensor const& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts) {
|
||||
CHECK_INPUT(output, "output must be a CUDA tensor");
|
||||
CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor");
|
||||
CHECK_INPUT(input, "input must be a CUDA tensor");
|
||||
CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor");
|
||||
CHECK_INPUT(input_offset_by_experts,
|
||||
"input_offset_by_experts must be a CUDA tensor");
|
||||
CHECK_INPUT(output_scale_offset_by_experts,
|
||||
"output_scale_offset_by_experts must be a CUDA tensor");
|
||||
torch::Tensor const& output_scale_offset_by_experts, int64_t m_topk,
|
||||
int64_t k) {
|
||||
CHECK_INPUT(output, "output");
|
||||
CHECK_INPUT(output_scale, "output_scale");
|
||||
CHECK_INPUT(input, "input");
|
||||
CHECK_INPUT(input_global_scale, "input_global_scale");
|
||||
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
|
||||
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
|
||||
|
||||
TORCH_CHECK(output.dim() == 2);
|
||||
TORCH_CHECK(output_scale.dim() == 2);
|
||||
@@ -335,8 +371,6 @@ void scaled_fp4_experts_quant_sm1xxa(
|
||||
TORCH_CHECK(output_scale.scalar_type() == INT);
|
||||
|
||||
const int BLOCK_SIZE = 16;
|
||||
auto m_topk = input.size(0);
|
||||
auto k = input.size(1);
|
||||
TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
|
||||
auto n_experts = input_global_scale.size(0);
|
||||
TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
|
||||
@@ -348,7 +382,21 @@ void scaled_fp4_experts_quant_sm1xxa(
|
||||
int padded_k = (scales_k + (4 - 1)) / 4 * 4;
|
||||
// 4 means 4 fp8 values are packed into one int32
|
||||
TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
|
||||
}
|
||||
|
||||
void scaled_fp4_experts_quant_sm1xxa(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts) {
|
||||
auto m_topk = input.size(0);
|
||||
auto k = input.size(1);
|
||||
|
||||
validate_fp4_experts_quant_inputs(output, output_scale, input,
|
||||
input_global_scale, input_offset_by_experts,
|
||||
output_scale_offset_by_experts, m_topk, k);
|
||||
|
||||
auto n_experts = input_global_scale.size(0);
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
const cudaStream_t stream =
|
||||
at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
@@ -356,7 +404,38 @@ void scaled_fp4_experts_quant_sm1xxa(
|
||||
VLLM_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
vllm::quant_impl<cuda_type>(
|
||||
vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
|
||||
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||
input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(),
|
||||
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts) {
|
||||
auto m_topk = input.size(0);
|
||||
// Input has gate || up layout, so k = input.size(1) / 2
|
||||
auto k_times_2 = input.size(1);
|
||||
TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
|
||||
auto k = k_times_2 / 2;
|
||||
|
||||
validate_fp4_experts_quant_inputs(output, output_scale, input,
|
||||
input_global_scale, input_offset_by_experts,
|
||||
output_scale_offset_by_experts, m_topk, k);
|
||||
|
||||
auto n_experts = input_global_scale.size(0);
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
const cudaStream_t stream =
|
||||
at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "silu_mul_nvfp4_experts_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
|
||||
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||
input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(),
|
||||
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||
|
||||
@@ -41,6 +41,15 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,
|
||||
torch::Tensor& input_sf);
|
||||
#endif
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts);
|
||||
#endif
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf) {
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
@@ -74,3 +83,18 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No compiled silu_and_mul nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
void silu_and_mul_scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts) {
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
output, output_scale, input, input_global_scale, input_offset_by_experts,
|
||||
output_scale_offset_by_experts);
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
|
||||
}
|
||||
|
||||
@@ -239,4 +239,34 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
||||
return e2m1Vec;
|
||||
}
|
||||
|
||||
// silu in float32
|
||||
__device__ __forceinline__ float silu(float x) {
|
||||
return __fdividef(x, (1.f + __expf(-x)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float2 silu2(float2 x) {
|
||||
return make_float2(silu(x.x), silu(x.y));
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__inline__ __device__ PackedVec<Type> compute_silu_mul(
|
||||
const PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
|
||||
PackedVec<Type> result;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
|
||||
// silu_mul in float32
|
||||
if constexpr (std::is_same_v<Type, half>) {
|
||||
float2 silu_vec = silu2(__half22float2(x_vec.elts[i]));
|
||||
result.elts[i] = __float22half2_rn(
|
||||
__fmul2_rn(silu_vec, __half22float2(y_vec.elts[i])));
|
||||
} else {
|
||||
float2 silu_vec = silu2(__bfloat1622float2(x_vec.elts[i]));
|
||||
result.elts[i] = __float22bfloat162_rn(
|
||||
__fmul2_rn(silu_vec, __bfloat1622float2(y_vec.elts[i])));
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
@@ -3,6 +3,8 @@
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
constexpr uint64_t THREADS_PER_EXPERT = 512;
|
||||
@@ -114,22 +116,17 @@ inline void launch_compute_problem_sizes(const torch::Tensor& topk_ids,
|
||||
const bool swap_ab) {
|
||||
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
|
||||
|
||||
const int32_t* topk_ptr = static_cast<const int32_t*>(topk_ids.data_ptr());
|
||||
int32_t* ps1_ptr = static_cast<int32_t*>(problem_sizes1.data_ptr());
|
||||
int32_t* ps2_ptr = static_cast<int32_t*>(problem_sizes2.data_ptr());
|
||||
int32_t* atomic_ptr = static_cast<int32_t*>(atomic_buffer.data_ptr());
|
||||
auto const* topk_ptr = topk_ids.data_ptr<int32_t>();
|
||||
auto* ps1_ptr = problem_sizes1.data_ptr<int32_t>();
|
||||
auto* ps2_ptr = problem_sizes2.data_ptr<int32_t>();
|
||||
auto* atomic_ptr = atomic_buffer.data_ptr<int32_t>();
|
||||
|
||||
if (swap_ab) {
|
||||
compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(
|
||||
VLLM_DISPATCH_BOOL(swap_ab, SwapAB, [&] {
|
||||
compute_problem_sizes<SwapAB><<<num_experts, num_threads, 0, stream>>>(
|
||||
topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,
|
||||
static_cast<int>(topk_ids.numel()), static_cast<int>(n),
|
||||
static_cast<int>(k));
|
||||
} else {
|
||||
compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(
|
||||
topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,
|
||||
static_cast<int>(topk_ids.numel()), static_cast<int>(n),
|
||||
static_cast<int>(k));
|
||||
}
|
||||
});
|
||||
}
|
||||
} // namespace
|
||||
|
||||
@@ -153,6 +150,93 @@ void get_cutlass_moe_mm_problem_sizes_caller(
|
||||
may_swap_ab);
|
||||
}
|
||||
|
||||
template <bool SWAP_AB>
|
||||
__global__ void compute_problem_sizes_from_expert_offsets(
|
||||
const int64_t* __restrict__ expert_first_token_offset,
|
||||
int32_t* __restrict__ problem_sizes1, int32_t* __restrict__ problem_sizes2,
|
||||
const int num_experts, const int n, const int k) {
|
||||
int const expert_id = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (expert_id >= num_experts) {
|
||||
return;
|
||||
}
|
||||
|
||||
int64_t const m64 = expert_first_token_offset[expert_id + 1] -
|
||||
expert_first_token_offset[expert_id];
|
||||
int32_t const m = static_cast<int32_t>(m64);
|
||||
|
||||
int32_t* ps1 = problem_sizes1 + expert_id * 3;
|
||||
int32_t* ps2 = problem_sizes2 + expert_id * 3;
|
||||
|
||||
if constexpr (!SWAP_AB) {
|
||||
// [M, 2*N, K]
|
||||
ps1[0] = m;
|
||||
ps1[1] = 2 * n;
|
||||
ps1[2] = k;
|
||||
// [M, K, N]
|
||||
ps2[0] = m;
|
||||
ps2[1] = k;
|
||||
ps2[2] = n;
|
||||
} else {
|
||||
// swap logical M/N in the problem shape
|
||||
// [2*N, M, K]
|
||||
ps1[0] = 2 * n;
|
||||
ps1[1] = m;
|
||||
ps1[2] = k;
|
||||
// [K, M, N]
|
||||
ps2[0] = k;
|
||||
ps2[1] = m;
|
||||
ps2[2] = n;
|
||||
}
|
||||
}
|
||||
|
||||
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
|
||||
const torch::Tensor& expert_first_token_offset,
|
||||
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
|
||||
const int64_t n, const int64_t k, const bool swap_ab) {
|
||||
TORCH_CHECK(expert_first_token_offset.is_cuda(),
|
||||
"expert_first_token_offset must be a CUDA tensor");
|
||||
TORCH_CHECK(expert_first_token_offset.dtype() == torch::kInt64,
|
||||
"expert_first_token_offset must be int64");
|
||||
|
||||
TORCH_CHECK(problem_sizes1.is_cuda() && problem_sizes2.is_cuda(),
|
||||
"problem_sizes must be CUDA tensors");
|
||||
TORCH_CHECK(problem_sizes1.dtype() == torch::kInt32 &&
|
||||
problem_sizes2.dtype() == torch::kInt32,
|
||||
"problem_sizes must be int32");
|
||||
TORCH_CHECK(problem_sizes1.is_contiguous() && problem_sizes2.is_contiguous(),
|
||||
"problem_sizes must be contiguous");
|
||||
TORCH_CHECK(problem_sizes1.dim() == 2 && problem_sizes2.dim() == 2,
|
||||
"problem_sizes must be 2D tensors");
|
||||
TORCH_CHECK(problem_sizes1.size(1) == 3 && problem_sizes2.size(1) == 3,
|
||||
"problem_sizes second dim must be 3");
|
||||
TORCH_CHECK(problem_sizes1.sizes() == problem_sizes2.sizes(),
|
||||
"problem_sizes1 and problem_sizes2 must have same shape");
|
||||
|
||||
int64_t const num_experts64 = problem_sizes1.size(0);
|
||||
TORCH_CHECK(expert_first_token_offset.numel() == num_experts64 + 1,
|
||||
"expert_first_token_offset must have num_experts + 1 elements");
|
||||
TORCH_CHECK(num_experts64 <= INT32_MAX, "num_experts must fit in int32");
|
||||
TORCH_CHECK(n <= INT32_MAX && k <= INT32_MAX, "n and k must fit in int32");
|
||||
|
||||
int const num_experts = static_cast<int>(num_experts64);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(
|
||||
expert_first_token_offset.device().index());
|
||||
|
||||
int const threads = (num_experts < 256) ? num_experts : 256;
|
||||
int const blocks = (num_experts + threads - 1) / threads;
|
||||
|
||||
auto const* offsets_ptr = expert_first_token_offset.data_ptr<int64_t>();
|
||||
auto* ps1_ptr = problem_sizes1.data_ptr<int32_t>();
|
||||
auto* ps2_ptr = problem_sizes2.data_ptr<int32_t>();
|
||||
|
||||
VLLM_DISPATCH_BOOL(swap_ab, SwapAB, [&] {
|
||||
compute_problem_sizes_from_expert_offsets<SwapAB>
|
||||
<<<blocks, threads, 0, stream>>>(offsets_ptr, ps1_ptr, ps2_ptr,
|
||||
num_experts, static_cast<int>(n),
|
||||
static_cast<int>(k));
|
||||
});
|
||||
}
|
||||
|
||||
void get_cutlass_moe_mm_data_caller(
|
||||
const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
|
||||
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
|
||||
|
||||
@@ -83,6 +83,11 @@ void get_cutlass_moe_mm_problem_sizes_caller(
|
||||
const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
|
||||
std::optional<bool> force_swap_ab = std::nullopt);
|
||||
|
||||
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
|
||||
const torch::Tensor& expert_first_token_offset,
|
||||
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
|
||||
const int64_t n, const int64_t k, const bool swap_ab);
|
||||
|
||||
void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
|
||||
torch::Tensor& problem_sizes1,
|
||||
torch::Tensor& problem_sizes2,
|
||||
@@ -322,6 +327,25 @@ void get_cutlass_moe_mm_problem_sizes(
|
||||
version_num, ". Required capability: 90, 100, or 120");
|
||||
}
|
||||
|
||||
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
|
||||
const torch::Tensor& expert_first_token_offset,
|
||||
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
|
||||
const int64_t n, const int64_t k, const bool swap_ab) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
|
||||
expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
|
||||
return;
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled get_cutlass_moe_mm_problem_sizes_from_expert_offsets: "
|
||||
"no cutlass_scaled_mm kernel for CUDA device capability: ",
|
||||
version_num, ". Required capability: 90, 100, or 120");
|
||||
}
|
||||
|
||||
void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
|
||||
torch::Tensor& problem_sizes1,
|
||||
torch::Tensor& problem_sizes2,
|
||||
|
||||
@@ -4,28 +4,77 @@
|
||||
#include "quantization/vectorization_utils.cuh"
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <ATen/cuda/Exceptions.h>
|
||||
#include <tuple>
|
||||
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, typename fp8_type>
|
||||
__global__ void scaled_fp8_quant_kernel_strided(
|
||||
// STRIDE_I_ZERO: true if scale_stride_i == 0 (per-tensor or per-channel)
|
||||
// STRIDE_J_ZERO: true if scale_stride_j == 0 (per-tensor or per-token)
|
||||
template <typename scalar_t, typename fp8_type, bool STRIDE_I_ZERO,
|
||||
bool STRIDE_J_ZERO>
|
||||
__global__ void scaled_fp8_quant_kernel_strided_group_shape(
|
||||
fp8_type* __restrict__ out, const scalar_t* __restrict__ input,
|
||||
const float* __restrict__ scale, int hidden_size, int64_t in_row_stride,
|
||||
int64_t out_row_stride) {
|
||||
const int64_t token_idx = blockIdx.x; // one token per block
|
||||
int64_t out_row_stride, int group_m, int group_n, int64_t scale_stride_i,
|
||||
int64_t scale_stride_j) {
|
||||
const int64_t token_idx = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
|
||||
const scalar_t* token_in = input + token_idx * in_row_stride;
|
||||
fp8_type* token_out = out + token_idx * out_row_stride;
|
||||
|
||||
const float inv_scale = 1.0f / (*scale);
|
||||
// Precompute row-level base offset for scale access (compile-time eliminated
|
||||
// when STRIDE_I_ZERO)
|
||||
const int64_t scale_row_base =
|
||||
STRIDE_I_ZERO ? 0
|
||||
: static_cast<int>(token_idx) / group_m * scale_stride_i;
|
||||
|
||||
vectorize_with_alignment<16>(
|
||||
token_in, token_out, hidden_size, tid, blockDim.x,
|
||||
[=] __device__(fp8_type & dst, const scalar_t& src) {
|
||||
dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src),
|
||||
inv_scale);
|
||||
});
|
||||
auto get_inv_scale = [&](int gj) {
|
||||
return 1.0f / scale[scale_row_base + gj * scale_stride_j];
|
||||
};
|
||||
|
||||
int cached_gj = -1;
|
||||
float cached_inv_scale = 0.0f;
|
||||
auto get_inv_scale_cached = [&](int gj) {
|
||||
if (gj != cached_gj) {
|
||||
cached_inv_scale = 1.0f / scale[scale_row_base + gj * scale_stride_j];
|
||||
cached_gj = gj;
|
||||
}
|
||||
return cached_inv_scale;
|
||||
};
|
||||
|
||||
constexpr int VEC_SIZE = 16; // FP8 so vectorize to 128 bits
|
||||
auto scaled_fp8_conversion_vectorized = [&](const scalar_t* in, fp8_type* out,
|
||||
int size, float inv_scale) {
|
||||
vectorize_with_alignment<VEC_SIZE>(
|
||||
in, out, size, tid, blockDim.x,
|
||||
[=] __device__(fp8_type & dst, const scalar_t& src) {
|
||||
dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src),
|
||||
inv_scale);
|
||||
});
|
||||
};
|
||||
|
||||
if (STRIDE_J_ZERO && hidden_size % VEC_SIZE == 0) {
|
||||
// Per-tensor or per-token: single scale per row, vectorize full row
|
||||
scaled_fp8_conversion_vectorized(token_in, token_out, hidden_size,
|
||||
get_inv_scale(0));
|
||||
} else if (group_n % VEC_SIZE == 0) {
|
||||
// Multiple column groups with vectorization
|
||||
const int num_groups_n = hidden_size / group_n;
|
||||
|
||||
for (int gj = 0; gj < num_groups_n; gj++) {
|
||||
scaled_fp8_conversion_vectorized(token_in + gj * group_n,
|
||||
token_out + gj * group_n, group_n,
|
||||
get_inv_scale(gj));
|
||||
}
|
||||
} else {
|
||||
// Scalar path for small column groups (group_n < VEC_SIZE)
|
||||
for (int n = tid; n < hidden_size; n += blockDim.x) {
|
||||
const int gj = n / group_n;
|
||||
token_out[n] = scaled_fp8_conversion<true, fp8_type>(
|
||||
static_cast<float>(token_in[n]), get_inv_scale_cached(gj));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename fp8_type>
|
||||
@@ -133,17 +182,116 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor const& input, // [..., d]
|
||||
torch::Tensor const& scale) // [1]
|
||||
void static_scaled_fp8_quant(
|
||||
torch::Tensor& out, // [..., d]
|
||||
torch::Tensor const& input, // [..., d]
|
||||
torch::Tensor const& scale, // various shapes
|
||||
std::optional<std::tuple<int64_t, int64_t>>
|
||||
opt_group_shape) // optional explicit (group_m, group_n)
|
||||
{
|
||||
TORCH_CHECK(input.stride(-1) == 1,
|
||||
"last dimension of input must be contiguous");
|
||||
TORCH_CHECK(out.stride(-1) == 1,
|
||||
"last dimension of output must be contiguous");
|
||||
|
||||
const int hidden_size = input.size(-1);
|
||||
const int num_tokens = input.numel() / hidden_size;
|
||||
const int hidden_size = input.size(-1); // N (columns)
|
||||
const int num_tokens = input.numel() / hidden_size; // M (rows)
|
||||
|
||||
// Determine group_m, group_n, and scale strides from scale shape
|
||||
// Scale indexing: scale[gi * scale_stride_j + gj * scale_stride_i]
|
||||
// where gi = m / group_m, gj = n / group_n
|
||||
int group_m, group_n;
|
||||
int64_t scale_stride_i, scale_stride_j;
|
||||
|
||||
if (scale.dim() == 0 || scale.numel() == 1) {
|
||||
// Per-tensor: one scale for the entire tensor
|
||||
group_m = num_tokens;
|
||||
group_n = hidden_size;
|
||||
scale_stride_i = 0;
|
||||
scale_stride_j = 0;
|
||||
} else if (scale.dim() == 1) {
|
||||
// 1D scale: require explicit group_shape to disambiguate per-channel vs
|
||||
// per-token (avoids edge case where num_tokens == hidden_size)
|
||||
TORCH_CHECK(opt_group_shape.has_value(),
|
||||
"1D scale requires explicit group_shape to disambiguate "
|
||||
"per-channel vs per-token quantization. "
|
||||
"Use group_shape=(-1, 1) for per-channel or group_shape=(1, "
|
||||
"-1) for per-token.");
|
||||
|
||||
const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
|
||||
group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
|
||||
group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
|
||||
|
||||
// Validate the explicit group shape matches the 1D scale
|
||||
const int64_t scale_len = scale.numel();
|
||||
const int64_t expected_scale_m = num_tokens / group_m;
|
||||
const int64_t expected_scale_n = hidden_size / group_n;
|
||||
const int64_t expected_scale_numel = expected_scale_m * expected_scale_n;
|
||||
|
||||
TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (",
|
||||
scale_len, ") does not match expected size (",
|
||||
expected_scale_numel, ") for group_shape (", opt_group_m, ", ",
|
||||
opt_group_n, ") with input shape (", num_tokens, ", ",
|
||||
hidden_size, ")");
|
||||
|
||||
// For 1D scale, determine strides based on which dim is trivial
|
||||
// Scale indexing: scale[gi * scale_stride_i + gj * scale_stride_j]
|
||||
// where gi = m / group_m (row group), gj = n / group_n (col group)
|
||||
if (expected_scale_m == 1) {
|
||||
// Per-channel style: one scale in M dim, scale varies along N
|
||||
// gi = 0 always, gj varies, so stride_1 traverses the scale
|
||||
scale_stride_i = 0;
|
||||
scale_stride_j = scale.stride(0);
|
||||
} else if (expected_scale_n == 1) {
|
||||
// Per-token style: one scale in N dim, scale varies along M
|
||||
// gj = 0 always, gi varies, so stride_0 traverses the scale
|
||||
scale_stride_i = scale.stride(0);
|
||||
scale_stride_j = 0;
|
||||
} else {
|
||||
TORCH_CHECK(
|
||||
false,
|
||||
"1D scale can only be used when one of the scale dimensions is 1. "
|
||||
"For 2D group scaling, use a 2D scale tensor.");
|
||||
}
|
||||
} else if (scale.dim() == 2) {
|
||||
// 2D scale: infer group sizes from scale dimensions (or use explicit if
|
||||
// provided)
|
||||
const int64_t scale_size_0 = scale.size(0);
|
||||
const int64_t scale_size_1 = scale.size(1);
|
||||
|
||||
TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens,
|
||||
") must be divisible by scale.size(0) (", scale_size_0, ")");
|
||||
TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (", hidden_size,
|
||||
") must be divisible by scale.size(1) (", scale_size_1, ")");
|
||||
|
||||
// Infer from 2D scale shape
|
||||
int inferred_group_m = num_tokens / scale_size_0;
|
||||
int inferred_group_n = hidden_size / scale_size_1;
|
||||
|
||||
// Use explicit if provided, otherwise use inferred
|
||||
if (opt_group_shape.has_value()) {
|
||||
const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
|
||||
group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
|
||||
group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
|
||||
|
||||
// Validate explicit matches inferred
|
||||
TORCH_CHECK(group_m == inferred_group_m && group_n == inferred_group_n,
|
||||
"Explicit group_shape (", opt_group_m, ", ", opt_group_n,
|
||||
") does not match inferred group shape (", inferred_group_m,
|
||||
", ", inferred_group_n, ") from 2D scale tensor shape (",
|
||||
scale_size_0, ", ", scale_size_1, ")");
|
||||
} else {
|
||||
group_m = inferred_group_m;
|
||||
group_n = inferred_group_n;
|
||||
}
|
||||
|
||||
scale_stride_i = scale.stride(0);
|
||||
scale_stride_j = scale.stride(1);
|
||||
} else {
|
||||
TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ",
|
||||
scale.dim(), "D");
|
||||
}
|
||||
|
||||
const int block_size = 256;
|
||||
dim3 grid(num_tokens);
|
||||
dim3 block(block_size);
|
||||
@@ -153,15 +301,23 @@ void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
// Dispatch to template-specialized kernel based on stride pattern
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
|
||||
VLLM_DISPATCH_FP8_TYPES(
|
||||
out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
|
||||
vllm::scaled_fp8_quant_kernel_strided<scalar_t, fp8_t>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
|
||||
scale.data_ptr<float>(), hidden_size, in_row_stride,
|
||||
out_row_stride);
|
||||
VLLM_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] {
|
||||
VLLM_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] {
|
||||
vllm::scaled_fp8_quant_kernel_strided_group_shape<
|
||||
scalar_t, fp8_t, S0_ZERO, S1_ZERO>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
|
||||
scale.data_ptr<float>(), hidden_size, in_row_stride,
|
||||
out_row_stride, group_m, group_n, scale_stride_i,
|
||||
scale_stride_j);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
+7
-4
@@ -1,3 +1,4 @@
|
||||
#include "cuda_compat.h"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
#include <torch/cuda.h>
|
||||
@@ -97,7 +98,9 @@ static inline __device__ bool isPartialMatch(float x, uint32_t pattern) {
|
||||
template <typename T, typename idxT, typename Func>
|
||||
__device__ void vectorized_process(size_t thread_rank, size_t num_threads,
|
||||
const T* in, idxT len, Func f) {
|
||||
constexpr int WARP_SIZE = 32;
|
||||
// Use dynamic WARP_SIZE from cuda_compat.h to support both
|
||||
// Wave64 (MI300X/gfx942) and Wave32 (Strix Halo/gfx1151) architectures
|
||||
constexpr int kWarpSize = WARP_SIZE;
|
||||
using WideT = float4;
|
||||
if constexpr (sizeof(T) >= sizeof(WideT)) {
|
||||
for (idxT i = thread_rank; i < len; i += num_threads) {
|
||||
@@ -132,8 +135,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
|
||||
}
|
||||
}
|
||||
|
||||
static_assert(WARP_SIZE >= items_per_scalar);
|
||||
// and because items_per_scalar > skip_cnt, WARP_SIZE > skip_cnt
|
||||
static_assert(kWarpSize >= items_per_scalar);
|
||||
// and because items_per_scalar > skip_cnt, kWarpSize > skip_cnt
|
||||
// no need to use loop
|
||||
if (thread_rank < skip_cnt) {
|
||||
f(in[thread_rank], thread_rank);
|
||||
@@ -142,7 +145,7 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
|
||||
// len_cast * items_per_scalar + items_per_scalar > len - skip_cnt;
|
||||
// and so
|
||||
// len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <=
|
||||
// WARP_SIZE no need to use loop
|
||||
// kWarpSize no need to use loop
|
||||
const idxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank;
|
||||
if (remain_i < len) {
|
||||
f(in[remain_i], remain_i);
|
||||
|
||||
+41
-2
@@ -487,6 +487,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,
|
||||
&get_cutlass_moe_mm_problem_sizes);
|
||||
|
||||
// compute per-expert problem sizes from expert_first_token_offset
|
||||
// produced by vLLM's moe_permute kernel
|
||||
ops.def(
|
||||
"get_cutlass_moe_mm_problem_sizes_from_expert_offsets("
|
||||
" Tensor expert_first_token_offset, "
|
||||
" Tensor! problem_sizes1, "
|
||||
" Tensor! problem_sizes2, "
|
||||
" int n, int k, bool swap_ab) -> ()");
|
||||
ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", torch::kCUDA,
|
||||
&get_cutlass_moe_mm_problem_sizes_from_expert_offsets);
|
||||
|
||||
// A function that computes data required to run fused MoE with w8a8 grouped
|
||||
// GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs
|
||||
// as an input, and computes expert_offsets (token start indices of each
|
||||
@@ -558,6 +569,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"Tensor output_scale_offset_by_experts) -> ()");
|
||||
ops.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);
|
||||
|
||||
// Fused SiLU+Mul+NVFP4 experts quantization.
|
||||
ops.def(
|
||||
"silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! "
|
||||
"output_scale,"
|
||||
"Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts) -> ()");
|
||||
ops.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA,
|
||||
&silu_and_mul_scaled_fp4_experts_quant);
|
||||
|
||||
// Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
|
||||
// of the given capability
|
||||
ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
|
||||
@@ -579,9 +599,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle);
|
||||
|
||||
// Compute FP8 quantized tensor for given scaling factor.
|
||||
// Supports per-tensor, per-channel, per-token, and arbitrary 2D group
|
||||
// scaling. Optional group_m/group_n specify the group shape explicitly;
|
||||
// required for 1D scales to disambiguate per-channel vs per-token.
|
||||
ops.def(
|
||||
"static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale) -> "
|
||||
"()");
|
||||
"static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
|
||||
"(int, int)? group_shape=None) -> ()");
|
||||
ops.impl("static_scaled_fp8_quant", torch::kCUDA, &static_scaled_fp8_quant);
|
||||
|
||||
// Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
|
||||
@@ -714,6 +737,22 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
|
||||
" Tensor scale) -> ()");
|
||||
cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
|
||||
|
||||
// Rotate Q and K, then write to kv cache for MLA
|
||||
cache_ops.def(
|
||||
"concat_and_cache_mla_rope_fused("
|
||||
" Tensor positions,"
|
||||
" Tensor! q_pe,"
|
||||
" Tensor! k_pe,"
|
||||
" Tensor kv_c,"
|
||||
" Tensor cos_sin_cache,"
|
||||
" bool is_neox,"
|
||||
" Tensor slot_mapping,"
|
||||
" Tensor! kv_cache,"
|
||||
" str kv_cache_dtype,"
|
||||
" Tensor kv_cache_scale) -> ()");
|
||||
cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
|
||||
&concat_and_cache_mla_rope_fused);
|
||||
|
||||
// Convert the key and value cache to fp8 data type.
|
||||
cache_ops.def(
|
||||
"convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
|
||||
|
||||
+25
-2
@@ -273,6 +273,7 @@ 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
|
||||
ARG PPLX_COMMIT_HASH
|
||||
ARG DEEPEP_COMMIT_HASH
|
||||
ARG NVSHMEM_VER
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p /tmp/ep_kernels_workspace/dist && \
|
||||
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
|
||||
@@ -280,7 +281,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--workspace /tmp/ep_kernels_workspace \
|
||||
--mode wheel \
|
||||
${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
|
||||
${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \
|
||||
${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \
|
||||
${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \
|
||||
find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
|
||||
#################### EXTENSIONS BUILD IMAGE ####################
|
||||
|
||||
@@ -615,6 +617,7 @@ RUN mv vllm src/vllm
|
||||
FROM vllm-base AS vllm-openai-base
|
||||
ARG TARGETPLATFORM
|
||||
ARG INSTALL_KV_CONNECTORS=false
|
||||
ARG CUDA_VERSION
|
||||
|
||||
ARG PIP_INDEX_URL UV_INDEX_URL
|
||||
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
@@ -624,10 +627,30 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# install kv_connectors if requested
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
|
||||
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
|
||||
CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \
|
||||
CUDA_HOME=/usr/local/cuda; \
|
||||
# lmcache requires explicit specifying CUDA_HOME
|
||||
BUILD_PKGS="libcusparse-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
libcusolver-dev-${CUDA_VERSION_DASH}"; \
|
||||
if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
|
||||
uv pip install --system -r /tmp/kv_connectors.txt || true; \
|
||||
if [ "$CUDA_MAJOR" -ge 13 ]; then \
|
||||
uv pip install --system nixl-cu13; \
|
||||
fi; \
|
||||
uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
|
||||
# if the above fails, install from source
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends ${BUILD_PKGS} && \
|
||||
uv pip install --system -r /tmp/kv_connectors.txt --no-build-isolation && \
|
||||
apt-get purge -y ${BUILD_PKGS} && \
|
||||
# clean up -dev packages, keep runtime libraries
|
||||
rm -rf /var/lib/apt/lists/* \
|
||||
); \
|
||||
fi
|
||||
|
||||
ENV VLLM_USAGE_SOURCE production-docker-image
|
||||
|
||||
@@ -22,13 +22,13 @@ RUN microdnf install -y dnf && dnf install -y gcc-toolset-14 make wget unzip \
|
||||
###############################################################
|
||||
FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS centos-deps-builder
|
||||
RUN microdnf install -y dnf && \
|
||||
dnf install -y https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-gpg-keys-9.0-24.el9.noarch.rpm \
|
||||
https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-stream-repos-9.0-24.el9.noarch.rpm \
|
||||
dnf install -y https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-gpg-keys-9.0-26.el9.noarch.rpm \
|
||||
https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-stream-repos-9.0-26.el9.noarch.rpm \
|
||||
https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \
|
||||
dnf config-manager --set-enabled crb
|
||||
|
||||
RUN dnf install -y openjpeg2-devel lcms2-devel tcl-devel tk-devel fribidi-devel && \
|
||||
dnf remove -y centos-gpg-keys-9.0-24.el9.noarch centos-stream-repos-9.0-24.el9.noarch
|
||||
RUN dnf install -y openjpeg2-devel lcms2-devel tcl-devel tk-devel fribidi-devel yajl-devel && \
|
||||
dnf remove -y centos-gpg-keys-9.0-24.el9.noarch centos-stream-repos-9.0-26.el9.noarch
|
||||
|
||||
|
||||
###############################################################
|
||||
@@ -346,4 +346,4 @@ WORKDIR /workspace/
|
||||
|
||||
RUN ln -s /workspace/vllm/tests && ln -s /workspace/vllm/examples && ln -s /workspace/vllm/benchmarks
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
@@ -11,6 +11,8 @@ ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="6af8b687"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="2d02c6a9"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
#TODO: When patch has been upstreamed, switch to the main repo/branch
|
||||
# ARG RIXL_BRANCH="<TODO>"
|
||||
@@ -31,6 +33,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
|
||||
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
|
||||
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
|
||||
ENV AITER_ROCM_ARCH=gfx942;gfx950
|
||||
ENV MORI_GPU_ARCHS=gfx942;gfx950
|
||||
|
||||
# Required for RCCL in ROCm7.1
|
||||
ENV HSA_NO_SCRATCH_RECLAIM=1
|
||||
@@ -44,7 +47,7 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install Python and other dependencies
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y software-properties-common git curl sudo vim less libgfortran5 \
|
||||
&& apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev \
|
||||
&& for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && break || \
|
||||
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
|
||||
@@ -119,6 +122,22 @@ RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
|
||||
&& cp /app/audio/dist/*.whl /app/install
|
||||
|
||||
|
||||
###
|
||||
### MORI Build
|
||||
###
|
||||
FROM base AS build_mori
|
||||
ARG MORI_BRANCH
|
||||
ARG MORI_REPO
|
||||
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
|
||||
pip install /install/*.whl
|
||||
RUN git clone ${MORI_REPO}
|
||||
RUN cd mori \
|
||||
&& git checkout ${MORI_BRANCH} \
|
||||
&& git submodule update --init --recursive \
|
||||
&& python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl
|
||||
RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
|
||||
|
||||
|
||||
###
|
||||
### RIXL Build
|
||||
###
|
||||
@@ -253,6 +272,8 @@ RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
|
||||
cp /install/*.whl /app/debs
|
||||
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
|
||||
cp /install/*.whl /app/debs
|
||||
RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
|
||||
cp /install/*.whl /app/debs
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install/,target=/install \
|
||||
cp /install/*.whl /app/debs
|
||||
|
||||
@@ -279,6 +300,8 @@ ARG ETCD_BRANCH
|
||||
ARG ETCD_REPO
|
||||
ARG UCX_BRANCH
|
||||
ARG UCX_REPO
|
||||
ARG MORI_BRANCH
|
||||
ARG MORI_REPO
|
||||
RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
|
||||
&& echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
|
||||
&& echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
|
||||
@@ -297,4 +320,6 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
|
||||
&& echo "ETCD_BRANCH: ${ETCD_BRANCH}" >> /app/versions.txt \
|
||||
&& echo "ETCD_REPO: ${ETCD_REPO}" >> /app/versions.txt \
|
||||
&& echo "UCX_BRANCH: ${UCX_BRANCH}" >> /app/versions.txt \
|
||||
&& echo "UCX_REPO: ${UCX_REPO}" >> /app/versions.txt
|
||||
&& echo "UCX_REPO: ${UCX_REPO}" >> /app/versions.txt \
|
||||
&& echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
|
||||
&& echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt
|
||||
|
||||
+1
-1
@@ -62,7 +62,7 @@ vLLM is flexible and easy to use with:
|
||||
|
||||
For more information, check out the following:
|
||||
|
||||
- [vLLM announcing blog post](https://vllm.ai) (intro to PagedAttention)
|
||||
- [vLLM announcing blog post](https://blog.vllm.ai/2023/06/20/vllm.html) (intro to PagedAttention)
|
||||
- [vLLM paper](https://arxiv.org/abs/2309.06180) (SOSP 2023)
|
||||
- [How continuous batching enables 23x throughput in LLM inference while reducing p50 latency](https://www.anyscale.com/blog/continuous-batching-llm-inference) by Cade Daniel et al.
|
||||
- [vLLM Meetups](community/meetups.md)
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
@@ -129,10 +129,10 @@ vllm bench sweep serve_sla \
|
||||
|
||||
The algorithm for adjusting the SLA variable is as follows:
|
||||
|
||||
1. Run the benchmark with infinite QPS, and use the corresponding metrics to determine the initial value of the variable.
|
||||
- For example, the initial request rate is set to the concurrency under infinite QPS.
|
||||
2. If the SLA is still satisfied, keep doubling the value until the SLA is no longer satisfied. This gives a relatively narrow window that contains the point where the SLA is barely satisfied.
|
||||
3. Apply binary search over the window to find the maximum value that still satisfies the SLA.
|
||||
1. Run the benchmark once with maximum possible QPS, and once with minimum possible QPS. For each run, calculate the distance of the SLA metrics from their targets, resulting in data points of QPS vs SLA distance.
|
||||
2. Perform spline interpolation between the data points to estimate the QPS that results in zero SLA distance.
|
||||
3. Run the benchmark with the estimated QPS and add the resulting data point to the history.
|
||||
4. Repeat Steps 2 and 3 until the maximum QPS that passes SLA and the minimum QPS that fails SLA in the history are close enough to each other.
|
||||
|
||||
!!! important
|
||||
SLA tuning is applied over each combination of `--serve-params`, `--bench-params`, and `--sla-params`.
|
||||
|
||||
@@ -142,7 +142,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.
|
||||
It is also worth noting that we should update `MAMBA_TYPE_TO_BACKEND_MAP` and `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/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.
|
||||
|
||||
@@ -149,7 +149,7 @@ The CUDA Graphs wrapper no longer manages the warm-up logic. The warm-up process
|
||||
|
||||
## CUDA Graphs Compatibility of Attention Backends
|
||||
|
||||
To signal the CUDA Graphs compatibility of the attention backends, we introduce a new enum type [AttentionCGSupport][vllm.v1.attention.backends.utils.AttentionCGSupport], which is an enum type that tracks the capability of the attention backend to support CUDA Graphs. The value is sorted in the order of the capability, i.e., `ALWAYS`> `UNIFORM_BATCH`> `UNIFORM_SINGLE_TOKEN_DECODE`> `NEVER`.
|
||||
To signal the CUDA Graphs compatibility of the attention backends, we introduce a new enum type [AttentionCGSupport][vllm.v1.attention.backend.AttentionCGSupport], which is an enum type that tracks the capability of the attention backend to support CUDA Graphs. The value is sorted in the order of the capability, i.e., `ALWAYS`> `UNIFORM_BATCH`> `UNIFORM_SINGLE_TOKEN_DECODE`> `NEVER`.
|
||||
|
||||
```python
|
||||
class AttentionCGSupport(enum.Enum):
|
||||
|
||||
@@ -0,0 +1,318 @@
|
||||
# CustomOp
|
||||
|
||||
`CustomOp` is an abstract class used for dispatching the forward method of various operations to the appropriate backend. It also offers a mechanism for both vLLM and OOT (Out-Of-Tree) plugins to register their custom operations.
|
||||
|
||||
This document will introduce how CustomOp works in vLLM and how to implement a new `CustomOp`.
|
||||
|
||||
## How CustomOp Works in vLLM
|
||||
|
||||
`CustomOp` manages two dictionaries of all custom ops (i.e., op classes, indexed by registered name) in its class, for vLLM and OOT plugins respectively.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
class CustomOp(nn.Module):
|
||||
|
||||
op_registry: dict[str, type["CustomOp"]] = {}
|
||||
op_registry_oot: dict[str, type["CustomOp"]] = {}
|
||||
```
|
||||
|
||||
We can use `@CustomOp.register("op_name")` to register an op class to the `CustomOp` system. After this, the `op_name` and its class will be added into the `op_registry` dictionary. In addition, We can also register an OOT op by `@CustomOp.register_oot("op_name")`. We will introduce this mechanism in detail later.
|
||||
|
||||
When a `CustomOp` is called (i.e., call its `forward()` method), if it is enabled (i.e., with `--compilation_config.custom_ops '["+op_name"]'`), it will automatically dispatch the forward method to the appropriate backend according to `current_platform`. Otherwise (i.e., it is disabled), it will only call the `forward_native()` method to use PyTorch-native implementation of this forward method.
|
||||
|
||||
- **CPU platform:** dispatch to `forward_cpu()`.
|
||||
- **CUDA platform:** dispatch to `forward_cuda()`.
|
||||
- **ROCm platform:** dispatch to `forward_hip()`. If `forward_hip()` is not implemented, it will use `forward_cuda()` as a fallback.
|
||||
- **XPU platform:** dispatch to `forward_xpu()`.
|
||||
- **TPU platform:** dispatch to `forward_tpu()`.
|
||||
- **OOT platform:** dispatch to `forward_oot()`. This will only be called on OOT platforms.
|
||||
- **Default:** dispatch to `forward_native()` as a final fallback for all platforms.
|
||||
|
||||
!!! note
|
||||
Note that the dispatching logic might not be absolute because of class inheritance. Derived class might override the behavior.
|
||||
|
||||
Furthermore, vLLM decides whether to enable or disable a `CustomOp` based on `compilation_config.custom_ops`. To be specific, if a `CustomOp` is not registered in `compilation_config.custom_ops` (i.e., uses the default config), it will be enabled if `compilation_config.custom_ops` contains `all`, or will be disabled if it contains `none`.
|
||||
|
||||
!!! note
|
||||
Note that `all` and `none` cannot coexist in `compilation_config.custom_ops`.
|
||||
|
||||
By default, if `compilation_config.backend == "inductor"` and `compilation_config.mode != CompilationMode.NONE`, a `none` will be appended into `compilation_config.custom_ops`, otherwise a `all` will be appended. In other words, this means `CustomOp` will be disabled in some platforms (i.e., those use `inductor` as dafault backend for `torch.compile`) when running with torch compile mode. In this case, Inductor generates (fused) Triton kernels for those disabled custom ops.
|
||||
|
||||
!!! note
|
||||
For multi-modal models, vLLM has enforece enabled some custom ops to use device-specific deep-optimized kernels for better performance in ViT part, such as `MMEncoderAttention` and `ApplyRotaryEmb`. We can also pass a `enforce_enable=True` param to the `__init__()` method of the `CustomOp` to enforce enable itself at object-level.
|
||||
|
||||
Note that this `enforce_enable` mechanism will be removed after we adding a separate `compilation_config` for multi-modal part.
|
||||
|
||||
## How to Customise Your Configuration for CustomOp
|
||||
|
||||
vLLM also offers fine-grained control over which custom ops to enable or disable for users, by manually passing a `--compilation_config.custom_ops '["..."]'` when launching a server.
|
||||
|
||||
For example:
|
||||
|
||||
- Use `--compilation_config.custom_ops '["all"]'` to enable all custom ops.
|
||||
- Use `--compilation_config.custom_ops '["none"]'` to disable all custom ops.
|
||||
- Use `--compilation_config.custom_ops '["all,-op1"]'` to enable all custom ops except op1 (i.e., prefixed with a `-` means "disable").
|
||||
- Use `--compilation_config.custom_ops '["none,+op1,+op2"]'` to only enable op1 and op2 (i.e., prefixed with a `+` means "enable").
|
||||
|
||||
## Types of Supported CustomOp in vLLM
|
||||
|
||||
**1. Attention:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"
|
||||
```
|
||||
|
||||
**2. Activation:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/activation.py:silu_and_mul"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:mul_and_silu"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:gelu_new"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:gelu_fast"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:quick_gelu"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:gelu_and_mul"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:gelu_and_mul_sparse"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:relu2"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:xielu"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:swigluoai_and_mul"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/activation.py:fatrelu_and_mul"
|
||||
```
|
||||
|
||||
**3. MM-Conv:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/conv.py:conv2d"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/conv.py:conv3d"
|
||||
```
|
||||
|
||||
**4. Embedding:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/vocab_parallel_embedding.py:vocab_parallel_embedding"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/vocab_parallel_embedding.py:parallel_lm_head"
|
||||
```
|
||||
|
||||
**5. Linear:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/linear.py:row_parallel_linear"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/linear.py:column_parallel_linear"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/linear.py:replicated_linear"
|
||||
```
|
||||
|
||||
**6. Logits Processor:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/logits_processor.py:logits_processor"
|
||||
```
|
||||
|
||||
**7. Mamba:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/mamba/mamba_mixer.py:mamba_mixer"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mamba_mixer2"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"
|
||||
|
||||
--8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
|
||||
```
|
||||
|
||||
**8. MoE:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/fused_moe/layer.py:fused_moe"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py:modular_fused_moe"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py:unquantized_fused_moe"
|
||||
|
||||
--8<-- "vllm/model_executor/models/transformers/moe.py:transformers_fused_moe"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/fused_moe/fused_moe.py:grouped_topk"
|
||||
```
|
||||
|
||||
**9. Norm:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/layernorm.py:rms_norm"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/layernorm.py:rms_norm_gated"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/layernorm.py:gemma_rms_norm"
|
||||
```
|
||||
|
||||
**10. Quantization:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/quantization/input_quant_fp8.py:quant_fp8"
|
||||
```
|
||||
|
||||
**11. Rope:**
|
||||
|
||||
```python
|
||||
--8<-- "vllm/model_executor/layers/rotary_embedding/base.py:rotary_embedding"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py:dual_chunk_rotary_embedding"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/rotary_embedding/common.py:apply_rotary_emb"
|
||||
```
|
||||
|
||||
## Guidelines for Implementing a New CustomOp
|
||||
|
||||
### Implement a New CustomOp in vLLM
|
||||
|
||||
This part is a tutorial of how to implement a New `CustomOp` in vLLM.
|
||||
|
||||
Steps:
|
||||
|
||||
1. Implement a new op class, which extends from `CustomOp` base class.
|
||||
2. Add the `@CustomOp.register("op_name")` decorator on this op class to register it into `CustomOp` system.
|
||||
3. Implement different `forward_xxx()` method according to your needs.
|
||||
|
||||
Taking `MMEncoderAttention` as an example:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
@CustomOp.register("mm_encoder_attn")
|
||||
class MMEncoderAttention(CustomOp):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
scale: float | None = None,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
multimodal_config: MultiModalConfig | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# Init...
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
|
||||
) -> torch.Tensor:
|
||||
# Call TORCH_SDPA implementation...
|
||||
|
||||
def forward_cuda(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
|
||||
) -> torch.Tensor:
|
||||
# Call FA or TORCH_SDPA implementation...
|
||||
|
||||
def forward_cpu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
|
||||
) -> torch.Tensor:
|
||||
# Call TORCH_SDPA implementation...
|
||||
|
||||
def forward_xpu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
|
||||
) -> torch.Tensor:
|
||||
# Call FA implementation...
|
||||
|
||||
def forward_tpu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
|
||||
) -> torch.Tensor:
|
||||
# Call PALLAS implementation...
|
||||
```
|
||||
|
||||
### Register a New CustomOp in OOT Device Plugins
|
||||
|
||||
Currently, thanks to [vLLM's hardware-plugin mechanism](./plugin_system.md), there are various OOT device plugins emerging out to enable vLLM seamlessly runs on different hardwares. You can also find more details about this mechanism at [Introducing vLLM Hardware Plugin, Best Practice from Ascend NPU](https://blog.vllm.ai/2025/05/12/hardware-plugin.html).
|
||||
|
||||
- **Official device plugins:** [vllm-ascend](https://github.com/vllm-project/vllm-ascend) (for Huawei Ascend NPU), [vllm-spyre](https://github.com/vllm-project/vllm-spyre)
|
||||
(for Spyre), [vllm-gaudi](https://github.com/vllm-project/vllm-gaudi) (for Intel Gaudi), [vllm-neuron](https://github.com/vllm-project/vllm-neuron) (for AWS Neuron), [vllm-meta](https://github.com/vllm-project/vllm-metal) (for Apple Silicon), etc.
|
||||
- **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), etc.
|
||||
|
||||
In this case, `CustomOp` can enable these hardware manufacturers to seamlessly replace vLLM's operations with their deep-optimized kernels for specific devices at runtime, by just registering an OOT `CustomOp` and implementing the `forward_oot()` method.
|
||||
|
||||
Now, this part will show you how to register an OOT `CustomOp` for a device plugin.
|
||||
|
||||
Taking `MMEncoderAttention` as an example:
|
||||
|
||||
1. Implement a `CustomMMEncoderAttention` class which extends from `MMEncoderAttention` and implement its `forward_oot()` method.
|
||||
2. Register your `CustomMMEncoderAttention` into vLLM to replace `MMEncoderAttention`.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
|
||||
from vllm.model_executor.custom_op import CustomOp
|
||||
|
||||
|
||||
@CustomOp.register_oot("MMEncoderAttention")
|
||||
class CustomMMEncoderAttention(MMEncoderAttention):
|
||||
|
||||
def __init__(...):
|
||||
super().__init__(...)
|
||||
|
||||
def forward_oot(...):
|
||||
# Call optimized device-specific kernels.
|
||||
...
|
||||
```
|
||||
|
||||
In this case, a new item `{"MMEncoderAttention": CustomMMEncoderAttention}` will be added into `op_registry_oot`. When initializing a `MMEncoderAttention` op object, if the class name (i.e., `MMEncoderAttention`) is contained in the keys of `op_registry_oot`, vLLM will replace it with our registered class (i.e., `CustomMMEncoderAttention`) and instantiate it.
|
||||
|
||||
After that, when this `MMEncoderAttention` op is called, your `forward_oot()` will be called if it is enabled. Thus, you will get expected performance on your hardwares without directly modify vLLM.
|
||||
|
||||
In addition, you can also register all your `CustomOp` at one place for better management.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm.model_executor.custom_op import CustomOp
|
||||
|
||||
|
||||
REGISTERED_CUSTOM_OPS = {
|
||||
"CustomOP1": YourCustomOp1,
|
||||
"CustomOP2": YourCustomOp2,
|
||||
"CustomOP3": YourCustomOp3,
|
||||
}
|
||||
|
||||
for op_name, op_cls in REGISTERED_CUSTOM_OPS.items():
|
||||
CustomOp.register_oot(_decorated_op_cls=op_cls, name=op_name)
|
||||
```
|
||||
@@ -86,13 +86,12 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
|
||||
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
|
||||
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
|
||||
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | [`deep_gemm_moe_fp8`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.deep_gemm_moe_fp8],</br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp4],</br>[`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`cutlass_moe_fp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp8],</br>[`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],</br>[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
|
||||
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
|
||||
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
|
||||
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | 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] |
|
||||
|
||||
@@ -139,18 +139,14 @@ token data.
|
||||
const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
|
||||
```
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/query.png" alt="query" width="70%" />
|
||||
</p>
|
||||

|
||||
|
||||
Each thread defines its own `q_ptr` which points to the assigned
|
||||
query token data on global memory. For example, if `VEC_SIZE` is 4
|
||||
and `HEAD_SIZE` is 128, the `q_ptr` points to data that contains
|
||||
total of 128 elements divided into 128 / 4 = 32 vecs.
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/q_vecs.png" alt="q_vecs" width="70%" />
|
||||
</p>
|
||||

|
||||
|
||||
```cpp
|
||||
__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
|
||||
@@ -187,9 +183,7 @@ key token at different iterations. As shown above, that `k_ptr`
|
||||
points to key token data based on `k_cache` at assigned block,
|
||||
assigned head and assigned token.
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/key.png" alt="key" width="70%" />
|
||||
</p>
|
||||

|
||||
|
||||
The diagram above illustrates the memory layout for key data. It
|
||||
assumes that the `BLOCK_SIZE` is 16, `HEAD_SIZE` is 128, `x` is
|
||||
@@ -202,9 +196,7 @@ iterations. Inside each rectangle, there are a total 32 vecs (128
|
||||
elements for one token) that will be processed by 2 threads (one
|
||||
thread group) separately.
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/k_vecs.png" alt="k_vecs" width="70%" />
|
||||
</p>
|
||||

|
||||
|
||||
```cpp
|
||||
K_vec k_vecs[NUM_VECS_PER_THREAD]
|
||||
@@ -361,17 +353,11 @@ later steps. Now, it should store the normalized softmax result of
|
||||
|
||||
## Value
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/value.png" alt="value" width="70%" />
|
||||
</p>
|
||||

|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/logits_vec.png" alt="logits_vec" width="50%" />
|
||||
</p>
|
||||

|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/design/paged_attention/v_vec.png" alt="v_vec" width="70%" />
|
||||
</p>
|
||||

|
||||
|
||||
Now we need to retrieve the value data and perform dot multiplication
|
||||
with `logits`. Unlike query and key, there is no thread group
|
||||
|
||||
@@ -124,7 +124,7 @@ Every plugin has three parts:
|
||||
|
||||
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.
|
||||
5. Implement the attention backend class `MyDummyAttention` in `my_dummy_attention.py`. The attention backend class should inherit from [AttentionBackend][vllm.v1.attention.backend.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:
|
||||
|
||||
@@ -153,5 +153,5 @@ The interface for the model/module may change during vLLM's development. If you
|
||||
|
||||
!!! warning "Deprecations"
|
||||
- `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
|
||||
- `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
|
||||
- `seed_everything` platform interface is deprecated. It will be removed in v0.14.0 or later. Please use `vllm.utils.torch_utils.set_random_seed` instead.
|
||||
- `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
|
||||
- `seed_everything` platform interface is deprecated. It will be removed in v0.15.0 or later. Please use `vllm.utils.torch_utils.set_random_seed` instead.
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# torch.compile with Multimodal Encoders
|
||||
|
||||
`torch.compile` can now be applied to multimodal encoders and miscellaneous nn modules in vLLM, including vision-language models like LLaMA 4, Qwen-VL,
|
||||
and similar encoder-based architectures.
|
||||
|
||||
This document covers the basics of how the `torch.compile` integration works for multimodal encoders in vLLM, as well as how to apply the decorator
|
||||
to new models to improve performance.
|
||||
|
||||
!!! note
|
||||
For general information about `torch.compile` integration in vLLM, see the [torch.compile design document](./torch_compile.md).
|
||||
|
||||
## Overview
|
||||
|
||||
We have recently enabled the `@supports_torch_compile` decorator to work for multiple nn module components within a model type; this enables
|
||||
turning compile on for multimodal encoders, bringing performance improvements to additional components of the stack.
|
||||
|
||||
When applied to the vision block of [`Qwen2_5_vl`](https://github.com/vllm-project/vllm/pull/23207) we observe ~4.5% e2e perf improvements with
|
||||
some increase in compilation time
|
||||
|
||||
This feature is off by default, but can be enabled by setting `compile_mm_encoder: true` in the compilation config when models have the
|
||||
`@supports_torch_compile` decorator.
|
||||
|
||||
## How Compilation Works for Multimodal Components
|
||||
|
||||
### APIs for Enablement
|
||||
|
||||
To compile a multimodal component such as an encoder, we follow the same mechanism as the LLM text backbone, with a few additional scaffoldings:
|
||||
|
||||
1. The `@supports_torch_compile` decorator should include `enable_if=should_torch_compile_mm_vit`. This will gate the compilation behind our
|
||||
`compile_mm_encoder` configuration
|
||||
|
||||
2. `with set_model_tag("<component_name>", is_encoder=True)` context manager should be used around the nn.Module's instantiation. Since torch.compile
|
||||
relies on caching artifacts to reduce start time, we must properly propagate the `<component_name>` information to the cache in order to avoid collisions
|
||||
with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
|
||||
components (see Compile Range Integration).
|
||||
|
||||
3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
|
||||
for torch.compile integration.
|
||||
|
||||
### CompilationConfig
|
||||
|
||||
With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend
|
||||
this for more configuration in the future.
|
||||
|
||||
## Applying torch.compile to a New Multimodal Model/Component
|
||||
|
||||
To apply `supports_torch_compile` to a new general nn.Module, we advise following the same steps in [`debug_vllm_compile`](./debug_vllm_compile.md); this includes:
|
||||
|
||||
1. Applying `supports_torch_compile` on initially small modules (such as basic MLP layers), then raising to more general modules until one reaches a good performance
|
||||
tradeoff
|
||||
|
||||
2. Leveraging [`tlparse`](https://github.com/meta-pytorch/tlparse) to identify and eliminate the source of recompiles and graph breaks
|
||||
|
||||
3. Using `dynamic_arg_dims` and proper `dynamic_shapes_config` to handle dynamism.
|
||||
|
||||
### Common pitfalls
|
||||
|
||||
## VllmBackend Feature Support
|
||||
|
||||
### Compile ranges
|
||||
|
||||
The torch.compile integration will try to rely on max_batch_size to infer compilation ranges for dynamic shapes; however, for modules used in the encoder, this
|
||||
shape can be difficult to infer due to the unspecified range of shapes the encoder may see as input. Therefore, we rely on `is_encoder=True` in the `set_model_tag`
|
||||
to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).
|
||||
|
||||
!!! note
|
||||
We may seek to tighten this range for better performance in the future
|
||||
|
||||
### Cudagraphs
|
||||
|
||||
We have not yet explored compilation for multimodal encoders with CUDAGraph integration; behavior is currently unspecified.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Graph Breaks in Vision Encoders
|
||||
|
||||
Some vision encoder operations may cause graph breaks. To identify them:
|
||||
|
||||
```bash
|
||||
TORCH_LOGS="+dynamo" vllm serve <MODEL>
|
||||
```
|
||||
|
||||
Common causes of graph breaks in multimodal models:
|
||||
|
||||
- **Dynamic image sizes**: Use `dynamic_shapes_config` to handle variable resolutions
|
||||
- **Untraceable operations**: Some operations (such as to_list) may not be supported by Dynamo
|
||||
- **Conditional processing**: Data-dependent branching based on image properties
|
||||
|
||||
### Compilation Errors
|
||||
|
||||
If compilation fails for a multimodal model:
|
||||
|
||||
1. **Disable and test**: First verify the model works without compilation:
|
||||
```bash
|
||||
VLLM_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'
|
||||
```
|
||||
|
||||
2. **Check logs**: Enable debug logging to see compilation details:
|
||||
```bash
|
||||
VLLM_LOGGING_LEVEL=DEBUG vllm serve <model> --compilation-config='{"compile_mm_encoder":"true"}'
|
||||
```
|
||||
|
||||
3. **Report issues**: If you find a bug, [open an issue on GitHub](https://github.com/vllm-project/vllm/issues/new/choose)
|
||||
|
||||
## See Also
|
||||
|
||||
- [torch.compile Integration](./torch_compile.md) - Core design document
|
||||
- [Debugging torch.compile](./debug_vllm_compile.md) - Detailed debugging guide
|
||||
- [Multimodal Inputs](../features/multimodal_inputs.md) - How to pass multimodal data
|
||||
- [Disaggregated Encoder](../features/disagg_encoder.md) - Scaling vision encoders
|
||||
- [Supported Multimodal Models](../models/supported_models.md#list-of-multimodal-language-models) - Model compatibility
|
||||
@@ -37,10 +37,10 @@ For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
|
||||
```
|
||||
|
||||
- **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and number of blocks to allocate (per worker):
|
||||
- **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and total CPU memory bytes to allocate:
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "num_cpu_blocks": 1000}}'
|
||||
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
|
||||
```
|
||||
|
||||
## Benchmarks
|
||||
|
||||
+13
-16
@@ -10,7 +10,7 @@ them locally with
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
sql_lora_path = snapshot_download(repo_id="yard1/llama-2-7b-sql-lora-test")
|
||||
sql_lora_path = snapshot_download(repo_id="jeeejeee/llama32-3b-text2sql-spider")
|
||||
```
|
||||
|
||||
Then we instantiate the base model and pass in the `enable_lora=True` flag:
|
||||
@@ -19,7 +19,7 @@ Then we instantiate the base model and pass in the `enable_lora=True` flag:
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.lora.request import LoRARequest
|
||||
|
||||
llm = LLM(model="meta-llama/Llama-2-7b-hf", enable_lora=True)
|
||||
llm = LLM(model="meta-llama/Llama-3.2-3B-Instruct", enable_lora=True)
|
||||
```
|
||||
|
||||
We can now submit the prompts and call `llm.generate` with the `lora_request` parameter. The first parameter
|
||||
@@ -55,14 +55,11 @@ LoRA adapted models can also be served with the Open-AI compatible vLLM server.
|
||||
`--lora-modules {name}={path} {name}={path}` to specify each LoRA module when we kick off the server:
|
||||
|
||||
```bash
|
||||
vllm serve meta-llama/Llama-2-7b-hf \
|
||||
vllm serve meta-llama/Llama-3.2-3B-Instruct \
|
||||
--enable-lora \
|
||||
--lora-modules sql-lora=$HOME/.cache/huggingface/hub/models--yard1--llama-2-7b-sql-lora-test/snapshots/0dfa347e8877a4d4ed19ee56c140fa518470028c/
|
||||
--lora-modules sql-lora=jeeejeee/llama32-3b-text2sql-spider
|
||||
```
|
||||
|
||||
!!! note
|
||||
The commit ID `0dfa347e8877a4d4ed19ee56c140fa518470028c` may change over time. Please check the latest commit ID in your environment to ensure you are using the correct one.
|
||||
|
||||
The server entrypoint accepts all other LoRA configuration parameters (`max_loras`, `max_lora_rank`, `max_cpu_loras`,
|
||||
etc.), which will apply to all forthcoming requests. Upon querying the `/models` endpoint, we should see our LoRA along
|
||||
with its base model (if `jq` is not installed, you can follow [this guide](https://jqlang.org/download/) to install it.):
|
||||
@@ -75,7 +72,7 @@ with its base model (if `jq` is not installed, you can follow [this guide](https
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "meta-llama/Llama-2-7b-hf",
|
||||
"id": "meta-llama/Llama-3.2-3B-Instruct",
|
||||
"object": "model",
|
||||
...
|
||||
},
|
||||
@@ -218,14 +215,14 @@ Alternatively, follow these example steps to implement your own plugin:
|
||||
In the previous version, users would provide LoRA modules via the following format, either as a key-value pair or in JSON format. For example:
|
||||
|
||||
```bash
|
||||
--lora-modules sql-lora=$HOME/.cache/huggingface/hub/models--yard1--llama-2-7b-sql-lora-test/snapshots/0dfa347e8877a4d4ed19ee56c140fa518470028c/
|
||||
--lora-modules sql-lora=jeeejeee/llama32-3b-text2sql-spider
|
||||
```
|
||||
|
||||
This would only include the `name` and `path` for each LoRA module, but did not provide a way to specify a `base_model_name`.
|
||||
Now, you can specify a base_model_name alongside the name and path using JSON format. For example:
|
||||
|
||||
```bash
|
||||
--lora-modules '{"name": "sql-lora", "path": "/path/to/lora", "base_model_name": "meta-llama/Llama-2-7b"}'
|
||||
--lora-modules '{"name": "sql-lora", "path": "jeeejeee/llama32-3b-text2sql-spider", "base_model_name": "meta-llama/Llama-3.2-3B-Instruct"}'
|
||||
```
|
||||
|
||||
To provide the backward compatibility support, you can still use the old key-value format (name=path), but the `base_model_name` will remain unspecified in that case.
|
||||
@@ -234,7 +231,7 @@ To provide the backward compatibility support, you can still use the old key-val
|
||||
|
||||
The new format of `--lora-modules` is mainly to support the display of parent model information in the model card. Here's an explanation of how your current response supports this:
|
||||
|
||||
- The `parent` field of LoRA model `sql-lora` now links to its base model `meta-llama/Llama-2-7b-hf`. This correctly reflects the hierarchical relationship between the base model and the LoRA adapter.
|
||||
- The `parent` field of LoRA model `sql-lora` now links to its base model `meta-llama/Llama-3.2-3B-Instruct`. This correctly reflects the hierarchical relationship between the base model and the LoRA adapter.
|
||||
- The `root` field points to the artifact location of the lora adapter.
|
||||
|
||||
??? console "Command output"
|
||||
@@ -246,11 +243,11 @@ The new format of `--lora-modules` is mainly to support the display of parent mo
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "meta-llama/Llama-2-7b-hf",
|
||||
"id": "meta-llama/Llama-3.2-3B-Instruct",
|
||||
"object": "model",
|
||||
"created": 1715644056,
|
||||
"owned_by": "vllm",
|
||||
"root": "~/.cache/huggingface/hub/models--meta-llama--Llama-2-7b-hf/snapshots/01c7f73d771dfac7d292323805ebc428287df4f9/",
|
||||
"root": "meta-llama/Llama-3.2-3B-Instruct",
|
||||
"parent": null,
|
||||
"permission": [
|
||||
{
|
||||
@@ -263,8 +260,8 @@ The new format of `--lora-modules` is mainly to support the display of parent mo
|
||||
"object": "model",
|
||||
"created": 1715644056,
|
||||
"owned_by": "vllm",
|
||||
"root": "~/.cache/huggingface/hub/models--yard1--llama-2-7b-sql-lora-test/snapshots/0dfa347e8877a4d4ed19ee56c140fa518470028c/",
|
||||
"parent": meta-llama/Llama-2-7b-hf,
|
||||
"root": "jeeejeee/llama32-3b-text2sql-spider",
|
||||
"parent": "meta-llama/Llama-3.2-3B-Instruct",
|
||||
"permission": [
|
||||
{
|
||||
....
|
||||
@@ -277,7 +274,7 @@ The new format of `--lora-modules` is mainly to support the display of parent mo
|
||||
|
||||
## LoRA Support for Tower and Connector of Multi-Modal Model
|
||||
|
||||
Currently, vLLM experimentally supports LoRA for the Tower and Connector components of multi-modal models. To enable this feature, you need to implement the corresponding token helper functions for the tower and connector. For more details on the rationale behind this approach, please refer to [PR 26674](https://github.com/vllm-project/vllm/pull/26674). We welcome contributions to extend LoRA support to additional models' tower and connector.
|
||||
Currently, vLLM experimentally supports LoRA for the Tower and Connector components of multi-modal models. To enable this feature, you need to implement the corresponding token helper functions for the tower and connector. For more details on the rationale behind this approach, please refer to [PR 26674](https://github.com/vllm-project/vllm/pull/26674). We welcome contributions to extend LoRA support to additional models' tower and connector. Please refer to [Issue 31479](https://github.com/vllm-project/vllm/issues/31479) to check the current model support status.
|
||||
|
||||
## Default LoRA Models For Multimodal Models
|
||||
|
||||
|
||||
@@ -166,49 +166,51 @@ Full example: [examples/offline_inference/vision_language_multi_image.py](../../
|
||||
|
||||
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.assets.image import ImageAsset
|
||||
??? code
|
||||
|
||||
llm = LLM(model="llava-hf/llava-1.5-7b-hf")
|
||||
image_url = "https://picsum.photos/id/32/512/512"
|
||||
image_pil = ImageAsset('cherry_blossom').pil_image
|
||||
image_embeds = torch.load(...)
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
conversation = [
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hello! How can I assist you today?"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_pil",
|
||||
"image_pil": image_pil,
|
||||
},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": image_embeds,
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in these images?",
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
llm = LLM(model="llava-hf/llava-1.5-7b-hf")
|
||||
image_url = "https://picsum.photos/id/32/512/512"
|
||||
image_pil = ImageAsset('cherry_blossom').pil_image
|
||||
image_embeds = torch.load(...)
|
||||
|
||||
# Perform inference and log output.
|
||||
outputs = llm.chat(conversation)
|
||||
conversation = [
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hello! How can I assist you today?"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_pil",
|
||||
"image_pil": image_pil,
|
||||
},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": image_embeds,
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in these images?",
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
```
|
||||
# Perform inference and log output.
|
||||
outputs = llm.chat(conversation)
|
||||
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Multi-image input can be extended to perform video captioning. We show this with [Qwen2-VL](https://huggingface.co/Qwen/Qwen2-VL-2B-Instruct) as it supports videos:
|
||||
|
||||
@@ -354,6 +356,44 @@ You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the mult
|
||||
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
|
||||
#### Automatic Audio Channel Normalization
|
||||
|
||||
vLLM automatically normalizes audio channels for models that require specific audio formats. When loading audio with libraries like `torchaudio`, stereo files return shape `[channels, time]`, but many audio models (particularly Whisper-based models) expect mono audio with shape `[time]`.
|
||||
|
||||
**Supported models with automatic mono conversion:**
|
||||
|
||||
- **Whisper** and all Whisper-based models
|
||||
- **Qwen2-Audio**
|
||||
- **Qwen2.5-Omni** / **Qwen3-Omni** (inherits from Qwen2.5-Omni)
|
||||
- **Ultravox**
|
||||
|
||||
For these models, vLLM automatically:
|
||||
|
||||
1. Detects if the model requires mono audio via the feature extractor
|
||||
2. Converts multi-channel audio to mono using channel averaging
|
||||
3. Handles both `(channels, time)` format (torchaudio) and `(time, channels)` format (soundfile)
|
||||
|
||||
**Example with stereo audio:**
|
||||
|
||||
```python
|
||||
import torchaudio
|
||||
from vllm import LLM
|
||||
|
||||
# Load stereo audio file - returns (channels, time) shape
|
||||
audio, sr = torchaudio.load("stereo_audio.wav")
|
||||
print(f"Original shape: {audio.shape}") # e.g., torch.Size([2, 16000])
|
||||
|
||||
# vLLM automatically converts to mono for Whisper-based models
|
||||
llm = LLM(model="openai/whisper-large-v3")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": "",
|
||||
"multi_modal_data": {"audio": (audio.numpy(), sr)},
|
||||
})
|
||||
```
|
||||
|
||||
No manual conversion is needed - vLLM handles the channel normalization automatically based on the model's requirements.
|
||||
|
||||
### Embedding Inputs
|
||||
|
||||
To input pre-computed embeddings belonging to a data type (i.e. image, video, or audio) directly to the language model,
|
||||
@@ -687,6 +727,31 @@ Full example: [examples/online_serving/openai_chat_completion_client_for_multimo
|
||||
export VLLM_VIDEO_FETCH_TIMEOUT=<timeout>
|
||||
```
|
||||
|
||||
#### Video Frame Recovery
|
||||
|
||||
For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place.
|
||||
|
||||
To enable video frame recovery, pass the `frame_recovery` parameter via `--media-io-kwargs`:
|
||||
|
||||
```bash
|
||||
# Example: Enable frame recovery
|
||||
vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
--media-io-kwargs '{"video": {"frame_recovery": true}}'
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `frame_recovery`: Boolean flag to enable forward-scan recovery. When `true`, failed frames are recovered using the next available frame within the dynamic window (up to the next target frame). Default is `false`.
|
||||
|
||||
**How it works:**
|
||||
|
||||
1. The system reads frames sequentially
|
||||
2. If a target frame fails to grab, it's marked as "failed"
|
||||
3. The next successfully grabbed frame (before reaching the next target) is used to recover the failed frame
|
||||
4. This approach handles both mid-video corruption and end-of-video truncation
|
||||
|
||||
Works with common video formats like MP4 when using OpenCV backends.
|
||||
|
||||
#### Custom RGBA Background Color
|
||||
|
||||
To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
|
||||
@@ -893,6 +958,8 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
|
||||
For Online Serving, you can also skip sending media if you expect cache hits with provided UUIDs. You can do so by sending media like this:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
# Image/video/audio URL:
|
||||
{
|
||||
|
||||
@@ -400,7 +400,7 @@ Flags: `--tool-call-parser functiongemma --chat-template examples/tool_chat_temp
|
||||
|
||||
Supported models:
|
||||
|
||||
* `Qwen/Qwen3-480B-A35B-Instruct`
|
||||
* `Qwen/Qwen3-Coder-480B-A35B-Instruct`
|
||||
* `Qwen/Qwen3-Coder-30B-A3B-Instruct`
|
||||
|
||||
Flags: `--tool-call-parser qwen3_xml`
|
||||
|
||||
@@ -98,9 +98,24 @@ Currently, there are no pre-built ROCm wheels.
|
||||
!!! note
|
||||
- You will need to config the `$AITER_BRANCH_OR_COMMIT` for your purpose.
|
||||
- The validated `$AITER_BRANCH_OR_COMMIT` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base).
|
||||
|
||||
|
||||
4. Build vLLM. For example, vLLM on ROCM 7.0 can be built with the following steps:
|
||||
|
||||
4. If you want to use MORI for EP or PD disaggregation, you can install [MORI](https://github.com/ROCm/mori) using the following steps:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ROCm/mori.git
|
||||
cd mori
|
||||
git checkout $MORI_BRANCH_OR_COMMIT
|
||||
git submodule sync; git submodule update --init --recursive
|
||||
MORI_GPU_ARCHS="gfx942;gfx950" python3 install .
|
||||
```
|
||||
|
||||
!!! note
|
||||
- You will need to config the `$MORI_BRANCH_OR_COMMIT` for your purpose.
|
||||
- The validated `$MORI_BRANCH_OR_COMMIT` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base).
|
||||
|
||||
|
||||
5. Build vLLM. For example, vLLM on ROCM 7.0 can be built with the following steps:
|
||||
|
||||
???+ console "Commands"
|
||||
|
||||
|
||||
@@ -363,7 +363,7 @@ th {
|
||||
| `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
|
||||
| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ |
|
||||
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
|
||||
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
|
||||
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
|
||||
| `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
|
||||
| `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
|
||||
| `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -371,10 +371,11 @@ th {
|
||||
| `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
|
||||
| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | | ✅︎ |
|
||||
| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
|
||||
| `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
|
||||
| `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. |✅︎| ✅︎ |
|
||||
| `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `ExaoneMoeCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B`, etc. | | |
|
||||
| `Exaone4ForCausalLM` | EXAONE-4 | `LGAI-EXAONE/EXAONE-4.0-32B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ |
|
||||
| `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ |
|
||||
@@ -388,7 +389,7 @@ th {
|
||||
| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
|
||||
| `GPT2LMHeadModel` | GPT-2 | `gpt2`, `gpt2-xl`, etc. | | ✅︎ |
|
||||
| `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ |
|
||||
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | ✅︎ |
|
||||
| `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ |
|
||||
| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
|
||||
@@ -399,12 +400,15 @@ th {
|
||||
| `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ |
|
||||
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
|
||||
| `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ |
|
||||
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
|
||||
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
|
||||
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
|
||||
| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
|
||||
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
|
||||
| `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
|
||||
| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ |
|
||||
| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
|
||||
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -459,6 +463,9 @@ th {
|
||||
| `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |
|
||||
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Grok2 requires `tokenizer.tok.json` with `tiktoken` installed. You can optionally override MoE router renormalization with `moe_router_renormalize`.
|
||||
|
||||
Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it!
|
||||
|
||||
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
@@ -636,29 +643,7 @@ See [this page](../features/multimodal_inputs.md) on how to pass multi-modal inp
|
||||
For hybrid-only models such as Llama-4, Step3 and Mistral-3, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (e.g, `--limit-mm-per-prompt '{"image":0}`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
|
||||
|
||||
!!! note
|
||||
vLLM currently only supports dynamic LoRA adapters on the language backbone of multimodal models.
|
||||
If you wish to use a model with LoRA in the multi-modal encoder,
|
||||
please merge the weights into the base model first before running it in vLLM like a regular model.
|
||||
|
||||
```python
|
||||
from peft import PeftConfig, PeftModel
|
||||
from transformers import AutoModelForImageTextToText, AutoProcessor
|
||||
|
||||
def merge_and_save(model_id: str, output_dir: str):
|
||||
base_model = AutoModelForImageTextToText.from_pretrained(model_id)
|
||||
lora_model = PeftModel.from_pretrained(
|
||||
base_model,
|
||||
model_id,
|
||||
config=PeftConfig.from_pretrained(model_id),
|
||||
)
|
||||
model = lora_model.merge_and_unload().to(dtype=base_model.dtype)
|
||||
model._hf_peft_config_loaded = False # Needed to save the merged model
|
||||
|
||||
processor = AutoProcessor.from_pretrained(model_id)
|
||||
|
||||
model.save_pretrained(output_dir)
|
||||
processor.save_pretrained(output_dir)
|
||||
```
|
||||
vLLM currently supports adding LoRA adapters to the language backbone for most multimodal models. Additionally, vLLM now experimentally supports adding LoRA to the tower and connector modules for some multimodal models. See [this page](../features/lora.md).
|
||||
|
||||
### Generative Models
|
||||
|
||||
@@ -695,10 +680,12 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `InternS1ForConditionalGeneration` | Intern-S1 | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
|
||||
| `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
|
||||
| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
|
||||
| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
|
||||
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
|
||||
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
|
||||
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
|
||||
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -716,10 +703,10 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
|
||||
| `Ovis2_5` | Ovis2.5 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.5-9B`, etc. | | |
|
||||
| `PaddleOCRVLForConditionalGeneration` | Paddle-OCR | T + I<sup>+</sup> | `PaddlePaddle/PaddleOCR-VL`, etc. | | |
|
||||
| `PaliGemmaForConditionalGeneration` | PaliGemma, PaliGemma 2 | T + I<sup>E</sup> | `google/paligemma-3b-pt-224`, `google/paligemma-3b-mix-224`, `google/paligemma2-3b-ft-docci-448`, etc. | | ✅︎ |
|
||||
| `PaliGemmaForConditionalGeneration` | PaliGemma, PaliGemma 2 | T + I<sup>E</sup> | `google/paligemma-3b-pt-224`, `google/paligemma-3b-mix-224`, `google/paligemma2-3b-ft-docci-448`, etc. | ✅︎ | ✅︎ |
|
||||
| `Phi3VForCausalLM` | Phi-3-Vision, Phi-3.5-Vision | T + I<sup>E+</sup> | `microsoft/Phi-3-vision-128k-instruct`, `microsoft/Phi-3.5-vision-instruct`, etc. | | ✅︎ |
|
||||
| `Phi4MMForCausalLM` | Phi-4-multimodal | T + I<sup>+</sup> / T + A<sup>+</sup> / I<sup>+</sup> + A<sup>+</sup> | `microsoft/Phi-4-multimodal-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `PixtralForConditionalGeneration` | Ministral 3 (Mistral format), Mistral 3 (Mistral format), Mistral Large 3 (Mistral format), Pixtral (Mistral format) | T + I<sup>+</sup> | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, `mistralai/Mistral-Large-3-675B-Instruct-2512` `mistralai/Pixtral-12B-2409` etc. | | ✅︎ |
|
||||
| `PixtralForConditionalGeneration` | Ministral 3 (Mistral format), Mistral 3 (Mistral format), Mistral Large 3 (Mistral format), Pixtral (Mistral format) | T + I<sup>+</sup> | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, `mistralai/Mistral-Large-3-675B-Instruct-2512` `mistralai/Pixtral-12B-2409` etc. | ✅︎ | ✅︎ |
|
||||
| `QwenVLForConditionalGeneration`<sup>^</sup> | Qwen-VL | T + I<sup>E+</sup> | `Qwen/Qwen-VL`, `Qwen/Qwen-VL-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen2AudioForConditionalGeneration` | Qwen2-Audio | T + A<sup>+</sup> | `Qwen/Qwen2-Audio-7B-Instruct` | | ✅︎ |
|
||||
| `Qwen2VLForConditionalGeneration` | QVQ, Qwen2-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/QVQ-72B-Preview`, `Qwen/Qwen2-VL-7B-Instruct`, `Qwen/Qwen2-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -767,9 +754,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.
|
||||
For more details, please see: <https://github.com/vllm-project/vllm/pull/4087#issuecomment-2250397630>
|
||||
|
||||
!!! note
|
||||
For Qwen2.5-Omni and Qwen3-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`) is currently work in progress and not yet supported.
|
||||
|
||||
#### Transcription
|
||||
|
||||
Speech2Text models trained specifically for Automatic Speech Recognition.
|
||||
@@ -803,6 +787,7 @@ The following table lists those that are tested in vLLM.
|
||||
| `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | |
|
||||
| `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ |
|
||||
| `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ |
|
||||
| `Qwen3VLForConditionalGeneration`<sup>C</sup> | Qwen3-VL | T + I + V | `Qwen/Qwen3-VL-Embedding-2B`, etc. | ✅︎ | ✅︎ |
|
||||
| `SiglipModel` | SigLIP, SigLIP2 | T / I | `google/siglip-base-patch16-224`, `google/siglip2-base-patch16-224` | | |
|
||||
| `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |
|
||||
|
||||
@@ -819,10 +804,18 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
|
||||
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
|--------------|--------|--------|-------------------|----------------------|---------------------------|
|
||||
| `JinaVLForSequenceClassification` | JinaVL-based | T + I<sup>E+</sup> | `jinaai/jina-reranker-m0`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3VLForSequenceClassification` | Qwen3-VL-Reranker | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-Reranker-2B`(see note), etc. | ✅︎ | ✅︎ |
|
||||
|
||||
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./pooling_models.md#model-conversion))
|
||||
\* Feature support is the same as that of the original model.
|
||||
|
||||
!!! note
|
||||
Similar to Qwen3-Reranker, you need to use the following `--hf_overrides` to load the official original `Qwen3-VL-Reranker`.
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-VL-Reranker-2B --hf_overrides '{"architectures": ["Qwen3VLForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}'
|
||||
```
|
||||
|
||||
## Model Support Policy
|
||||
|
||||
At vLLM, we are committed to facilitating the integration and support of third-party models within our ecosystem. Our approach is designed to balance the need for robustness and the practical limitations of supporting a wide range of models. Here’s how we manage third-party model support:
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# Claude Code
|
||||
|
||||
[Claude Code](https://code.claude.com/docs/en/quickstart) is Anthropic's official agentic coding tool that lives in your terminal. It can understand your codebase, edit files, run commands, and help you write code more efficiently.
|
||||
|
||||
By pointing Claude Code at a vLLM server, you can use your own models as the backend instead of the Anthropic API. This is useful for:
|
||||
|
||||
- Running fully local/private coding assistance
|
||||
- Using open-weight models with tool calling capabilities
|
||||
- Testing and developing with custom models
|
||||
|
||||
## How It Works
|
||||
|
||||
vLLM implements the Anthropic Messages API, which is the same API that Claude Code uses to communicate with Anthropic's servers. By setting `ANTHROPIC_BASE_URL` to point at your vLLM server, Claude Code sends its requests to vLLM instead of Anthropic. vLLM then translates these requests to work with your local model and returns responses in the format Claude Code expects.
|
||||
|
||||
This means any model served by vLLM with proper tool calling support can act as a drop-in replacement for Claude models in Claude Code.
|
||||
|
||||
## Requirements
|
||||
|
||||
Claude Code requires a model with strong tool calling capabilities. The model must support the OpenAI-compatible tool calling API. See [Tool Calling](../../features/tool_calling.md) for details on enabling tool calling for your model.
|
||||
|
||||
## Installation
|
||||
|
||||
First, install Claude Code by following the [official installation guide](https://docs.anthropic.com/en/docs/claude-code/getting-started).
|
||||
|
||||
## Starting the vLLM Server
|
||||
|
||||
Start vLLM with a tool-calling capable model - here's an example using `openai/gpt-oss-120b`:
|
||||
|
||||
```bash
|
||||
vllm serve openai/gpt-oss-120b --served-model-name my-model --enable-auto-tool-choice --tool-call-parser openai
|
||||
```
|
||||
|
||||
For other models, you'll need to enable tool calling explicitly with `--enable-auto-tool-choice` and the right `--tool-call-parser`. Refer to the [Tool Calling documentation](../../features/tool_calling.md) for the correct flags for your model.
|
||||
|
||||
## Configuring Claude Code
|
||||
|
||||
Launch Claude Code with environment variables pointing to your vLLM server:
|
||||
|
||||
```bash
|
||||
ANTHROPIC_BASE_URL=http://localhost:8000 \
|
||||
ANTHROPIC_API_KEY=dummy \
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL=my-model \
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL=my-model \
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL=my-model \
|
||||
claude
|
||||
```
|
||||
|
||||
The environment variables:
|
||||
|
||||
| Variable | Description |
|
||||
| -------------------------------- | --------------------------------------------------------------------- |
|
||||
| `ANTHROPIC_BASE_URL` | Points to your vLLM server (default port is 8000) |
|
||||
| `ANTHROPIC_API_KEY` | Can be any value since vLLM doesn't require authentication by default |
|
||||
| `ANTHROPIC_DEFAULT_OPUS_MODEL` | Model name for Opus-tier requests |
|
||||
| `ANTHROPIC_DEFAULT_SONNET_MODEL` | Model name for Sonnet-tier requests |
|
||||
| `ANTHROPIC_DEFAULT_HAIKU_MODEL` | Model name for Haiku-tier requests |
|
||||
|
||||
!!! tip
|
||||
You can add these environment variables to your shell profile (e.g., `.bashrc`, `.zshrc`), Claude Code configuration file (`~/.claude/settings.json`), or create a wrapper script for convenience.
|
||||
|
||||
## Testing the Setup
|
||||
|
||||
Once Claude Code launches, try a simple prompt to verify the connection:
|
||||
|
||||

|
||||
|
||||
If the model responds correctly, your setup is working. You can now use Claude Code with your vLLM-served model for coding tasks.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**Connection refused**: Ensure vLLM is running and accessible at the specified URL. Check that the port matches.
|
||||
|
||||
**Tool calls not working**: Verify that your model supports tool calling and that you've enabled it with the correct `--tool-call-parser` flag. See [Tool Calling](../../features/tool_calling.md).
|
||||
|
||||
**Model not found**: Ensure the `--served-model-name` matches the model names in your environment variables. You cannot use model names with `/` in them, such as `openai/gpt-oss-120b` directly from Huggingface, so beware of that limitation with Claude Code.
|
||||
@@ -173,6 +173,14 @@ with `--enable-request-id-headers`.
|
||||
print(completion._request_id)
|
||||
```
|
||||
|
||||
## Offline API Documentation
|
||||
|
||||
The FastAPI `/docs` endpoint requires an internet connection by default. To enable offline access in air-gapped environments, use the `--enable-offline-docs` flag:
|
||||
|
||||
```bash
|
||||
vllm serve NousResearch/Meta-Llama-3-8B-Instruct --enable-offline-docs
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Completions API
|
||||
@@ -362,7 +370,7 @@ and passing a list of `messages` in the request. Refer to the examples below for
|
||||
`MrLight/dse-qwen2-2b-mrl-v1` requires a placeholder image of the minimum image size for text query embeddings. See the full code
|
||||
example below for details.
|
||||
|
||||
Full example: [examples/pooling/embed/openai_chat_embedding_client_for_multimodal.py](../../examples/pooling/embed/openai_chat_embedding_client_for_multimodal.py)
|
||||
Full example: [examples/pooling/embed/vision_embedding_online.py](../../examples/pooling/embed/vision_embedding_online.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -667,7 +675,7 @@ Usually, the score for a sentence pair refers to the similarity between two sent
|
||||
|
||||
You can find the documentation for cross encoder models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).
|
||||
|
||||
Code example: [examples/pooling/score/openai_cross_encoder_score.py](../../examples/pooling/score/openai_cross_encoder_score.py)
|
||||
Code example: [examples/pooling/score/score_api_online.py](../../examples/pooling/score/score_api_online.py)
|
||||
|
||||
#### Score Template
|
||||
|
||||
@@ -863,7 +871,10 @@ You can pass multi-modal inputs to scoring models by passing `content` including
|
||||
print("Scoring output:", response_json["data"][0]["score"])
|
||||
print("Scoring output:", response_json["data"][1]["score"])
|
||||
```
|
||||
Full example: [examples/pooling/score/openai_cross_encoder_score_for_multimodal.py](../../examples/pooling/score/openai_cross_encoder_score_for_multimodal.py)
|
||||
Full example:
|
||||
|
||||
- [examples/pooling/score/vision_score_api_online.py](../../examples/pooling/score/vision_score_api_online.py)
|
||||
- examples/pooling/score/vision_rerank_api_online.py](../../examples/pooling/score/vision_rerank_api_online.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -893,7 +904,7 @@ endpoints are compatible with both [Jina AI's re-rank API interface](https://jin
|
||||
[Cohere's re-rank API interface](https://docs.cohere.com/v2/reference/rerank) to ensure compatibility with
|
||||
popular open-source tools.
|
||||
|
||||
Code example: [examples/pooling/score/openai_reranker.py](../../examples/pooling/score/openai_reranker.py)
|
||||
Code example: [examples/pooling/score/rerank_api_online.py](../../examples/pooling/score/rerank_api_online.py)
|
||||
|
||||
#### Example Request
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -44,12 +45,12 @@ def main(args: dict):
|
||||
if top_k is not None:
|
||||
sampling_params.top_k = top_k
|
||||
|
||||
def print_outputs(outputs):
|
||||
def print_outputs(outputs: list[RequestOutput], prompts: list):
|
||||
assert len(outputs) == len(prompts)
|
||||
print("\nGenerated Outputs:\n" + "-" * 80)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
for i, output in enumerate(outputs):
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\n")
|
||||
print(f"Prompt: {prompts[i]!r}\n")
|
||||
print(f"Generated text: {generated_text!r}")
|
||||
print("-" * 80)
|
||||
|
||||
@@ -66,14 +67,19 @@ def main(args: dict):
|
||||
},
|
||||
]
|
||||
outputs = llm.chat(conversation, sampling_params, use_tqdm=False)
|
||||
print_outputs(outputs)
|
||||
print_outputs(
|
||||
outputs,
|
||||
[
|
||||
conversation,
|
||||
],
|
||||
)
|
||||
|
||||
# You can run batch inference with llm.chat API
|
||||
conversations = [conversation for _ in range(10)]
|
||||
|
||||
# We turn on tqdm progress bar to verify it's indeed running batch inference
|
||||
outputs = llm.chat(conversations, sampling_params, use_tqdm=True)
|
||||
print_outputs(outputs)
|
||||
print_outputs(outputs, conversations)
|
||||
|
||||
# A chat template can be optionally supplied.
|
||||
# If not, the model will use its default chat template.
|
||||
@@ -87,7 +93,7 @@ def main(args: dict):
|
||||
use_tqdm=False,
|
||||
chat_template=chat_template,
|
||||
)
|
||||
print_outputs(outputs)
|
||||
print_outputs(outputs, conversations)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -4,9 +4,6 @@
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.config import AttentionConfig
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -23,11 +20,6 @@ def parse_args():
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
if current_platform.is_rocm():
|
||||
args.attention_config = AttentionConfig(
|
||||
backend=AttentionBackendEnum.FLEX_ATTENTION
|
||||
)
|
||||
|
||||
# Sample prompts.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
|
||||
@@ -4,9 +4,6 @@
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.config import AttentionConfig
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -23,11 +20,6 @@ def parse_args():
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
if current_platform.is_rocm():
|
||||
args.attention_config = AttentionConfig(
|
||||
backend=AttentionBackendEnum.FLEX_ATTENTION
|
||||
)
|
||||
|
||||
# Sample prompts.
|
||||
text_1 = "What is the capital of France?"
|
||||
texts_2 = [
|
||||
|
||||
@@ -9,7 +9,7 @@ Usage:
|
||||
python examples/offline_inference/context_extension.py
|
||||
"""
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm import LLM, RequestOutput, SamplingParams
|
||||
|
||||
|
||||
def create_llm():
|
||||
@@ -45,13 +45,15 @@ def run_llm_chat(llm):
|
||||
{"role": "assistant", "content": "Hello! How can I assist you today?"},
|
||||
]
|
||||
outputs = llm.chat(conversation, sampling_params, use_tqdm=False)
|
||||
return outputs
|
||||
return outputs, [
|
||||
conversation,
|
||||
]
|
||||
|
||||
|
||||
def print_outputs(outputs):
|
||||
def print_outputs(outputs: list[RequestOutput], conversations: list):
|
||||
print("\nGenerated Outputs:\n" + "-" * 80)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
for i, output in enumerate(outputs):
|
||||
prompt = conversations[i]
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\n")
|
||||
print(f"Generated text: {generated_text!r}")
|
||||
@@ -60,8 +62,8 @@ def print_outputs(outputs):
|
||||
|
||||
def main():
|
||||
llm = create_llm()
|
||||
outputs = run_llm_chat(llm)
|
||||
print_outputs(outputs)
|
||||
outputs, conversations = run_llm_chat(llm)
|
||||
print_outputs(outputs, conversations)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -10,7 +10,6 @@ python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
-q mixed_modalities
|
||||
|
||||
# Read vision and audio inputs from a single video file
|
||||
# NOTE: V1 engine does not support interleaved modalities yet.
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
-q use_audio_in_video
|
||||
|
||||
|
||||
@@ -152,9 +152,12 @@ def main(args):
|
||||
|
||||
# print the generated text
|
||||
if args.print_output:
|
||||
for output in outputs:
|
||||
for i, output in enumerate(outputs):
|
||||
print("-" * 50)
|
||||
print(f"prompt: {output.prompt}")
|
||||
if not args.custom_mm_prompts:
|
||||
print(f"prompt: {prompts[i].prompt}")
|
||||
else:
|
||||
print(f"prompt: {prompts[i]}")
|
||||
print(f"generated text: {output.outputs[0].text}")
|
||||
print("-" * 50)
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ from dataclasses import asdict
|
||||
from typing import NamedTuple
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoTokenizer
|
||||
from transformers import AutoProcessor, AutoTokenizer
|
||||
|
||||
from vllm import LLM, EngineArgs, SamplingParams
|
||||
from vllm.assets.image import ImageAsset
|
||||
@@ -769,6 +769,33 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# Kanana-V
|
||||
def run_kanana_v(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
|
||||
model_name = "kakaocorp/kanana-1.5-v-3b-instruct"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=8192,
|
||||
trust_remote_code=True,
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
messages = [
|
||||
[{"role": "user", "content": f"<image>\n{question}"}] for question in questions
|
||||
]
|
||||
prompts = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# Keye-VL
|
||||
def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
model_name = "Kwai-Keye/Keye-VL-8B-Preview"
|
||||
@@ -875,6 +902,37 @@ def run_lightonocr(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
def run_lfm2_vl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
|
||||
model_name = "LiquidAI/LFM2-VL-450M"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=4096,
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
)
|
||||
|
||||
processor = AutoProcessor.from_pretrained(model_name)
|
||||
messages = [
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image"}, {"type": "text", "text": question}],
|
||||
}
|
||||
]
|
||||
for question in questions
|
||||
]
|
||||
prompts = processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
def run_llama4(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
|
||||
@@ -1845,10 +1903,12 @@ model_example_map = {
|
||||
"idefics3": run_idefics3,
|
||||
"interns1": run_interns1,
|
||||
"internvl_chat": run_internvl,
|
||||
"kanana_v": run_kanana_v,
|
||||
"keye_vl": run_keye_vl,
|
||||
"keye_vl1_5": run_keye_vl1_5,
|
||||
"kimi_vl": run_kimi_vl,
|
||||
"lightonocr": run_lightonocr,
|
||||
"lfm2_vl": run_lfm2_vl,
|
||||
"llama4": run_llama4,
|
||||
"llava": run_llava,
|
||||
"llava-next": run_llava_next,
|
||||
|
||||
@@ -0,0 +1,320 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import threading
|
||||
import uuid
|
||||
|
||||
import aiohttp
|
||||
import msgpack
|
||||
import zmq
|
||||
from quart import Quart, make_response, request
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
prefill_instances: list[dict] = []
|
||||
decode_instances: list[dict] = []
|
||||
request_nums = 0
|
||||
app = Quart(__name__)
|
||||
|
||||
IP_PORT_PATTERN = re.compile(r"//(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):(\d+)")
|
||||
|
||||
|
||||
TRANSFER_TYPE = None
|
||||
|
||||
|
||||
def _append_whole_dict_unique(target_list, data_dict):
|
||||
new_filtered = {k: v for k, v in data_dict.items() if k != "index"}
|
||||
for existed in target_list:
|
||||
existed_filtered = {k: v for k, v in existed.items() if k != "index"}
|
||||
if existed_filtered == new_filtered:
|
||||
return False
|
||||
print("!!APPEND!!", data_dict)
|
||||
target_list.append(data_dict)
|
||||
transfer_mode = data_dict.get("transfer_mode", "unknown")
|
||||
global TRANSFER_TYPE
|
||||
|
||||
if TRANSFER_TYPE is None:
|
||||
TRANSFER_TYPE = transfer_mode
|
||||
logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
|
||||
elif transfer_mode != TRANSFER_TYPE:
|
||||
raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
_list_lock = threading.RLock()
|
||||
|
||||
|
||||
def _listen_for_register(hostname, port):
|
||||
context = zmq.Context()
|
||||
router_socket = context.socket(zmq.ROUTER)
|
||||
router_socket.bind(f"tcp://{hostname}:{port}")
|
||||
poller = zmq.Poller()
|
||||
poller.register(router_socket, zmq.POLLIN)
|
||||
global prefill_instances
|
||||
global decode_instances
|
||||
|
||||
while True:
|
||||
socks = dict(poller.poll())
|
||||
if router_socket in socks:
|
||||
remote_addr, msg = router_socket.recv_multipart()
|
||||
data = msgpack.loads(msg)
|
||||
if data["type"] == "HELLO":
|
||||
pass
|
||||
elif (
|
||||
data["type"] == "register"
|
||||
and data["role"] == "P"
|
||||
and data["request_address"] not in prefill_instances
|
||||
):
|
||||
with _list_lock:
|
||||
_append_whole_dict_unique(prefill_instances, data)
|
||||
|
||||
elif (
|
||||
data["type"] == "register"
|
||||
and data["role"] == "D"
|
||||
and data["request_address"] not in decode_instances
|
||||
):
|
||||
with _list_lock:
|
||||
_append_whole_dict_unique(decode_instances, data)
|
||||
|
||||
|
||||
def start_service_discovery(hostname, port):
|
||||
if not hostname:
|
||||
hostname = socket.gethostname()
|
||||
if port == 0:
|
||||
raise ValueError("Port cannot be 0")
|
||||
|
||||
_listener_thread = threading.Thread(
|
||||
target=_listen_for_register, args=(hostname, port), daemon=True
|
||||
)
|
||||
_listener_thread.start()
|
||||
return _listener_thread
|
||||
|
||||
|
||||
async def send_request_to_prefill(
|
||||
endpoint, req_data, request_id, d_endpoint, dip, dport, selected_prefill_dp_rank
|
||||
):
|
||||
req_data_copy = req_data
|
||||
|
||||
req_data_copy["kv_transfer_params"].update(
|
||||
{
|
||||
"do_remote_decode": True,
|
||||
"do_remote_prefill": False,
|
||||
"remote_handshake_port": d_endpoint["handshake_port"],
|
||||
"remote_notify_port": d_endpoint["notify_port"],
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": dip,
|
||||
"remote_port": dport,
|
||||
}
|
||||
)
|
||||
req_data_copy["stream"] = False
|
||||
req_data_copy["max_tokens"] = 1
|
||||
if "max_completion_tokens" in req_data_copy:
|
||||
req_data_copy["max_completion_tokens"] = 1
|
||||
if "stream_options" in req_data_copy:
|
||||
del req_data_copy["stream_options"]
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
|
||||
) as session:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
|
||||
"X-Request-Id": request_id,
|
||||
}
|
||||
if selected_prefill_dp_rank is not None:
|
||||
headers["X-data-parallel-rank"] = str(selected_prefill_dp_rank)
|
||||
async with session.post(
|
||||
url=endpoint, json=req_data_copy, headers=headers
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
return await response.json()
|
||||
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"send_request_to_prefill response.status != 200response.status = ",
|
||||
response.status,
|
||||
)
|
||||
|
||||
|
||||
async def start_decode_request(endpoint, req_data, request_id):
|
||||
session = aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
|
||||
)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
|
||||
"X-Request-Id": request_id,
|
||||
}
|
||||
response = await session.post(url=endpoint, json=req_data, headers=headers)
|
||||
return session, response
|
||||
|
||||
|
||||
async def stream_decode_response(session, response, request_id):
|
||||
try:
|
||||
if response.status == 200:
|
||||
async for chunk_bytes in response.content.iter_chunked(1024):
|
||||
yield chunk_bytes
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"decode response.status != 200, status = {response.status}"
|
||||
)
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
|
||||
async def send_request_to_decode(endpoint, req_data, request_id):
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
|
||||
) as session:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
|
||||
"X-Request-Id": request_id,
|
||||
}
|
||||
async with session.post(
|
||||
url=endpoint, json=req_data, headers=headers
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
async for chunk_bytes in response.content.iter_chunked(1024):
|
||||
yield chunk_bytes
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"send_request_to_decode response.status != 200,response.statuus = ",
|
||||
response.status,
|
||||
)
|
||||
|
||||
|
||||
def example_round_robin_dp_loader(request_number, dp_size):
|
||||
return request_nums % dp_size
|
||||
|
||||
|
||||
@app.route("/v1/completions", methods=["POST"])
|
||||
@app.route("/v1/chat/completions", methods=["POST"])
|
||||
async def handle_request():
|
||||
try:
|
||||
with _list_lock:
|
||||
global request_nums
|
||||
request_nums += 1
|
||||
|
||||
def extract_ip_port_fast(url):
|
||||
match = IP_PORT_PATTERN.search(url)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid URL format: {url}")
|
||||
return match.groups()
|
||||
|
||||
req_data = await request.get_json()
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
prefill_instance_endpoint = None
|
||||
decode_instance_endpoint = None
|
||||
error_msg = (
|
||||
"Service Unavailable: No prefill or decode instances are registered."
|
||||
)
|
||||
if not prefill_instances or not decode_instances:
|
||||
return await make_response(
|
||||
(
|
||||
error_msg,
|
||||
503,
|
||||
)
|
||||
)
|
||||
pid = request_nums % len(prefill_instances)
|
||||
did = request_nums % len(decode_instances)
|
||||
prefill_instance_endpoint = prefill_instances[pid]
|
||||
decode_instance_endpoint = decode_instances[did]
|
||||
|
||||
selected_prefill_dp_rank = None
|
||||
if prefill_instance_endpoint["dp_size"] > 1:
|
||||
selected_prefill_dp_rank = example_round_robin_dp_loader(
|
||||
request_nums // len(prefill_instance_endpoint),
|
||||
prefill_instance_endpoint["dp_size"],
|
||||
)
|
||||
|
||||
dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
|
||||
ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])
|
||||
|
||||
req_data_to_prefill = copy.deepcopy(req_data)
|
||||
req_data_to_prefill["kv_transfer_params"] = {}
|
||||
req_data["kv_transfer_params"] = {}
|
||||
req_data_to_prefill["kv_transfer_params"]["remote_dp_size"] = (
|
||||
decode_instance_endpoint["dp_size"]
|
||||
)
|
||||
req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (
|
||||
decode_instance_endpoint["tp_size"]
|
||||
)
|
||||
|
||||
send_prefill_task = asyncio.create_task(
|
||||
send_request_to_prefill(
|
||||
prefill_instance_endpoint["request_address"],
|
||||
req_data_to_prefill,
|
||||
request_id,
|
||||
decode_instance_endpoint,
|
||||
dip,
|
||||
dport,
|
||||
selected_prefill_dp_rank,
|
||||
)
|
||||
)
|
||||
ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])
|
||||
|
||||
req_data["max_tokens"] -= 1
|
||||
|
||||
req_data["kv_transfer_params"] = {
|
||||
"do_remote_decode": False,
|
||||
"do_remote_prefill": True,
|
||||
"remote_handshake_port": prefill_instance_endpoint["handshake_port"],
|
||||
"remote_notify_port": prefill_instance_endpoint["notify_port"],
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": ip,
|
||||
"remote_port": port,
|
||||
}
|
||||
if TRANSFER_TYPE == "READ":
|
||||
# In read mode, prefill and decode are executed serially.
|
||||
prefill_response = await send_prefill_task
|
||||
req_data["kv_transfer_params"]["remote_engine_id"] = prefill_response[
|
||||
"kv_transfer_params"
|
||||
]["remote_engine_id"]
|
||||
req_data["kv_transfer_params"]["remote_block_ids"] = prefill_response[
|
||||
"kv_transfer_params"
|
||||
]["remote_block_ids"]
|
||||
|
||||
req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
|
||||
"dp_size"
|
||||
]
|
||||
req_data["kv_transfer_params"]["remote_tp_size"] = prefill_instance_endpoint[
|
||||
"tp_size"
|
||||
]
|
||||
|
||||
if selected_prefill_dp_rank is not None:
|
||||
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
|
||||
|
||||
decode_request_task = asyncio.create_task(
|
||||
start_decode_request(
|
||||
decode_instance_endpoint["request_address"], req_data, request_id
|
||||
)
|
||||
)
|
||||
|
||||
session, decode_response = await decode_request_task
|
||||
stream_generator = stream_decode_response(session, decode_response, request_id)
|
||||
response = await make_response(stream_generator)
|
||||
return response
|
||||
except Exception as e:
|
||||
logger.exception("An error occurred while handling the request: %s", e)
|
||||
return await make_response(
|
||||
(
|
||||
f"Internal Server Error: {e!s}",
|
||||
500,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = start_service_discovery("0.0.0.0", 36367)
|
||||
app.debug = True
|
||||
app.config["BODY_TIMEOUT"] = 360000
|
||||
app.config["RESPONSE_TIMEOUT"] = 360000
|
||||
|
||||
app.run(host="0.0.0.0", port=10001)
|
||||
t.join()
|
||||
@@ -9,11 +9,11 @@ from openai import OpenAI
|
||||
from vllm.assets.audio import AudioAsset
|
||||
|
||||
|
||||
def sync_openai(audio_path: str, client: OpenAI):
|
||||
def sync_openai(audio_path: str, client: OpenAI, model: str):
|
||||
with open(audio_path, "rb") as f:
|
||||
translation = client.audio.translations.create(
|
||||
file=f,
|
||||
model="openai/whisper-large-v3",
|
||||
model=model,
|
||||
response_format="json",
|
||||
temperature=0.0,
|
||||
# Additional params not provided by OpenAI API.
|
||||
@@ -26,11 +26,13 @@ def sync_openai(audio_path: str, client: OpenAI):
|
||||
print("translation result:", translation.text)
|
||||
|
||||
|
||||
async def stream_openai_response(audio_path: str, base_url: str, api_key: str):
|
||||
async def stream_openai_response(
|
||||
audio_path: str, base_url: str, api_key: str, model: str
|
||||
):
|
||||
data = {
|
||||
"language": "it",
|
||||
"stream": True,
|
||||
"model": "openai/whisper-large-v3",
|
||||
"model": model,
|
||||
}
|
||||
url = base_url + "/audio/translations"
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
@@ -66,9 +68,13 @@ def main():
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
sync_openai(foscolo, client)
|
||||
|
||||
model = client.models.list().data[0].id
|
||||
print(f"Using model: {model}")
|
||||
|
||||
sync_openai(foscolo, client, model)
|
||||
# Run the asynchronous function
|
||||
asyncio.run(stream_openai_response(foscolo, openai_api_base, openai_api_key))
|
||||
asyncio.run(stream_openai_response(foscolo, openai_api_base, openai_api_key, model))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference with
|
||||
the correct prompt format on vision language models for multimodal embedding.
|
||||
|
||||
For most models, the prompt format should follow corresponding examples
|
||||
on HuggingFace model repository.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from dataclasses import asdict
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
text = "A cat standing in the snow."
|
||||
multi_modal_data = {"image": fetch_image(image_url)}
|
||||
|
||||
|
||||
def print_embeddings(embeds):
|
||||
embeds_trimmed = (str(embeds[:4])[:-1] + ", ...]") if len(embeds) > 4 else embeds
|
||||
print(f"Embeddings: {embeds_trimmed} (size={len(embeds)})")
|
||||
|
||||
|
||||
def run_qwen3_vl():
|
||||
engine_args = EngineArgs(
|
||||
model="Qwen/Qwen3-VL-Embedding-2B",
|
||||
runner="pooling",
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
default_instruction = "Represent the user's input."
|
||||
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
text_prompt = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
|
||||
image_prompt = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}<|im_end|>\n<|im_start|>assistant\n"
|
||||
image_text_prompt = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}{text}<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
llm = LLM(**asdict(engine_args))
|
||||
|
||||
print("Text embedding output:")
|
||||
outputs = llm.embed(text_prompt, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": image_prompt,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": image_text_prompt,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
model_example_map = {
|
||||
"qwen3_vl": run_qwen3_vl,
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
"Script to run a specified VLM through vLLM offline api."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
choices=model_example_map.keys(),
|
||||
required=True,
|
||||
help="The name of the embedding model.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
model_example_map[args.model]()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
+130
-5
@@ -21,7 +21,8 @@ from PIL import Image
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
text = "A cat standing in the snow."
|
||||
|
||||
|
||||
def create_chat_embeddings(
|
||||
@@ -30,6 +31,8 @@ def create_chat_embeddings(
|
||||
messages: list[ChatCompletionMessageParam],
|
||||
model: str,
|
||||
encoding_format: Literal["base64", "float"] | NotGiven = NOT_GIVEN,
|
||||
continue_final_message: bool = False,
|
||||
add_special_tokens: bool = False,
|
||||
) -> CreateEmbeddingResponse:
|
||||
"""
|
||||
Convenience function for accessing vLLM's Chat Embeddings API,
|
||||
@@ -38,10 +41,21 @@ def create_chat_embeddings(
|
||||
return client.post(
|
||||
"/embeddings",
|
||||
cast_to=CreateEmbeddingResponse,
|
||||
body={"messages": messages, "model": model, "encoding_format": encoding_format},
|
||||
body={
|
||||
"messages": messages,
|
||||
"model": model,
|
||||
"encoding_format": encoding_format,
|
||||
"continue_final_message": continue_final_message,
|
||||
"add_special_tokens": add_special_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def print_embeddings(embeds):
|
||||
embeds_trimmed = (str(embeds[:4])[:-1] + ", ...]") if len(embeds) > 4 else embeds
|
||||
print(f"Embeddings: {embeds_trimmed} (size={len(embeds)})")
|
||||
|
||||
|
||||
def run_clip(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
@@ -145,6 +159,113 @@ def run_dse_qwen2_vl(client: OpenAI, model: str):
|
||||
print("Text embedding output:", response.data[0].embedding)
|
||||
|
||||
|
||||
def run_qwen3_vl(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-Embedding-2B \
|
||||
--runner pooling \
|
||||
--max-model-len 8192
|
||||
"""
|
||||
|
||||
default_instruction = "Represent the user's input."
|
||||
|
||||
print("Text embedding output:")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": default_instruction},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": text},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
continue_final_message=True,
|
||||
add_special_tokens=True,
|
||||
)
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": default_instruction},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
continue_final_message=True,
|
||||
add_special_tokens=True,
|
||||
)
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": default_instruction},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{text}",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
continue_final_message=True,
|
||||
add_special_tokens=True,
|
||||
)
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
|
||||
def run_siglip(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
@@ -213,7 +334,8 @@ def run_vlm2vec(client: OpenAI, model: str):
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image embedding output:", response.data[0].embedding)
|
||||
print("Image embedding output:")
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
@@ -233,7 +355,8 @@ def run_vlm2vec(client: OpenAI, model: str):
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image+Text embedding output:", response.data[0].embedding)
|
||||
print("Image+Text embedding output:")
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
@@ -249,11 +372,13 @@ def run_vlm2vec(client: OpenAI, model: str):
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Text embedding output:", response.data[0].embedding)
|
||||
print("Text embedding output:")
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
|
||||
model_example_map = {
|
||||
"clip": run_clip,
|
||||
"qwen3_vl": run_qwen3_vl,
|
||||
"dse_qwen2_vl": run_dse_qwen2_vl,
|
||||
"siglip": run_siglip,
|
||||
"vlm2vec": run_vlm2vec,
|
||||
@@ -133,6 +133,36 @@ def run_jinavl_reranker(query: Query) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
def run_qwen3_vl(query: Query) -> ModelRequestData:
|
||||
image_placeholder = "<vision_start><|image_pad|><vision_end>"
|
||||
if query["modality"] == "text":
|
||||
prompt = query["text"]
|
||||
image = None
|
||||
elif query["modality"] == "image":
|
||||
prompt = image_placeholder
|
||||
image = query["image"]
|
||||
elif query["modality"] == "text+image":
|
||||
text = query["text"]
|
||||
prompt = f"{image_placeholder}\n{text}"
|
||||
image = query["image"]
|
||||
else:
|
||||
modality = query["modality"]
|
||||
raise ValueError(f"Unsupported query modality: '{modality}'")
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model="Qwen/Qwen3-VL-Embedding-2B",
|
||||
runner="pooling",
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompt=prompt,
|
||||
image=image,
|
||||
)
|
||||
|
||||
|
||||
def run_siglip(query: Query) -> ModelRequestData:
|
||||
if query["modality"] == "text":
|
||||
prompt = query["text"]
|
||||
@@ -353,6 +383,7 @@ model_example_map = {
|
||||
"clip": run_clip,
|
||||
"e5_v": run_e5_v,
|
||||
"jinavl_reranker": run_jinavl_reranker,
|
||||
"qwen3_vl": run_qwen3_vl,
|
||||
"siglip": run_siglip,
|
||||
"vlm2vec_phi3v": run_vlm2vec_phi3v,
|
||||
"vlm2vec_qwen2vl": run_vlm2vec_qwen2vl,
|
||||
|
||||
@@ -108,7 +108,7 @@ method_map = {
|
||||
|
||||
|
||||
def converting(
|
||||
model_name, classifier_from_tokens, path, method, use_pad_token=False, device="cpu"
|
||||
model_name, classifier_from_tokens, path, method, use_sep_token=False, device="cpu"
|
||||
):
|
||||
"""
|
||||
Main conversion function to transform a CausalLM model to SequenceClassification.
|
||||
@@ -118,7 +118,7 @@ def converting(
|
||||
classifier_from_tokens: List of tokens used for classification
|
||||
path: Output path to save the converted model
|
||||
method: Conversion method ('from_2_way_softmax' or 'no_post_processing')
|
||||
use_pad_token: Whether to use padding token in the sequence classification model
|
||||
use_sep_token: Whether to use separating token in the sequence classification model
|
||||
device: Device to load the model on ('cpu' or 'cuda')
|
||||
"""
|
||||
assert method in method_map, f"Unknown method: {method}"
|
||||
@@ -149,10 +149,10 @@ def converting(
|
||||
causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
|
||||
)
|
||||
|
||||
# Configure padding token settings
|
||||
# Note: Reranker models typically don't use padding tokens by default
|
||||
seq_cls_model.config.use_pad_token = use_pad_token
|
||||
seq_cls_model.config.pad_token_id = tokenizer.pad_token_id
|
||||
# Configure separating token settings
|
||||
# Note: `llm as reranker` defaults to not using separating token.
|
||||
seq_cls_model.config.use_sep_token = use_sep_token
|
||||
seq_cls_model.config.sep_token_id = tokenizer.sep_token_id
|
||||
|
||||
# Save the converted model and tokenizer
|
||||
seq_cls_model.save_pretrained(path)
|
||||
@@ -203,6 +203,6 @@ if __name__ == "__main__":
|
||||
model_name=args.model_name,
|
||||
classifier_from_tokens=json.loads(args.classifier_from_tokens),
|
||||
method=args.method,
|
||||
use_pad_token=args.use_pad_token,
|
||||
use_sep_token=args.use_sep_token,
|
||||
path=args.path,
|
||||
)
|
||||
|
||||
@@ -1,60 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example online usage of Score API.
|
||||
|
||||
Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model", type=str, default="jinaai/jina-reranker-m0")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
api_url = f"http://{args.host}:{args.port}/score"
|
||||
model_name = args.model
|
||||
|
||||
text_1 = "slm markdown"
|
||||
text_2 = {
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}
|
||||
score_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
print("\nPrompt when text_1 is string and text_2 is a image list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,23 @@
|
||||
<|im_start|>system
|
||||
Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>
|
||||
<|im_start|>user
|
||||
<Instruct>: {{
|
||||
messages
|
||||
| selectattr("role", "eq", "system")
|
||||
| map(attribute="content")
|
||||
| first
|
||||
| default("Given a search query, retrieve relevant candidates that answer the query.")
|
||||
}}<Query>:{{
|
||||
messages
|
||||
| selectattr("role", "eq", "query")
|
||||
| map(attribute="content")
|
||||
| first
|
||||
}}
|
||||
<Document>:{{
|
||||
messages
|
||||
| selectattr("role", "eq", "document")
|
||||
| map(attribute="content")
|
||||
| first
|
||||
}}<|im_end|>
|
||||
<|im_start|>assistant
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example Python client for multimodal rerank API which is compatible with
|
||||
Jina and Cohere https://jina.ai/reranker
|
||||
|
||||
Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
|
||||
e.g.
|
||||
vllm serve jinaai/jina-reranker-m0 --runner pooling
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-Reranker-2B \
|
||||
--runner pooling \
|
||||
--max-model-len 4096 \
|
||||
--hf_overrides '{"architectures": ["Qwen3VLForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' \
|
||||
--chat-template examples/pooling/score/template/qwen3_vl_reranker.jinja
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
documents = {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, " # noqa: E501
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust." # noqa: E501
|
||||
),
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
rerank_url = base_url + "/rerank"
|
||||
|
||||
response = requests.get(models_url, headers=headers)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
data = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
}
|
||||
response = requests.post(rerank_url, headers=headers, json=data)
|
||||
|
||||
# Check the response
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,172 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference with
|
||||
vision language reranker models for multimodal scoring tasks.
|
||||
|
||||
Vision language rerankers score the relevance between a text query and
|
||||
multimodal documents (text + images/videos).
|
||||
"""
|
||||
|
||||
from argparse import Namespace
|
||||
from collections.abc import Callable
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.entrypoints.score_utils import ScoreMultiModalParam
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
TEMPLATE_HOME = Path(__file__).parent / "template"
|
||||
|
||||
|
||||
class RerankModelData(NamedTuple):
|
||||
engine_args: EngineArgs
|
||||
chat_template: str | None = None
|
||||
|
||||
|
||||
def run_jinavl_reranker(modality: str) -> RerankModelData:
|
||||
assert modality == "image"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model="jinaai/jina-reranker-m0",
|
||||
runner="pooling",
|
||||
max_model_len=32768,
|
||||
trust_remote_code=True,
|
||||
mm_processor_kwargs={
|
||||
"min_pixels": 3136,
|
||||
"max_pixels": 602112,
|
||||
},
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
)
|
||||
return RerankModelData(
|
||||
engine_args=engine_args,
|
||||
)
|
||||
|
||||
|
||||
def run_qwen3_vl_reranker(modality: str) -> RerankModelData:
|
||||
engine_args = EngineArgs(
|
||||
model="Qwen/Qwen3-VL-Reranker-2B",
|
||||
runner="pooling",
|
||||
max_model_len=16384,
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
# HuggingFace model configuration overrides required for compatibility
|
||||
hf_overrides={
|
||||
# Manually route to sequence classification architecture
|
||||
# This tells vLLM to use Qwen3VLForSequenceClassification instead of
|
||||
# the default Qwen3VLForConditionalGeneration
|
||||
"architectures": ["Qwen3VLForSequenceClassification"],
|
||||
# Specify which token logits to extract from the language model head
|
||||
# The original reranker uses "no" and "yes" token logits for scoring
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
# Enable special handling for original Qwen3-Reranker models
|
||||
# This flag triggers conversion logic that transforms the two token
|
||||
# vectors into a single classification vector
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
)
|
||||
chat_template_path = "qwen3_vl_reranker.jinja"
|
||||
chat_template = (TEMPLATE_HOME / chat_template_path).read_text()
|
||||
return RerankModelData(
|
||||
engine_args=engine_args,
|
||||
chat_template=chat_template,
|
||||
)
|
||||
|
||||
|
||||
model_example_map: dict[str, Callable[[str], RerankModelData]] = {
|
||||
"jinavl_reranker": run_jinavl_reranker,
|
||||
"qwen3_vl_reranker": run_qwen3_vl_reranker,
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Demo on using vLLM for offline inference with "
|
||||
"vision language reranker models for multimodal scoring tasks."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-name",
|
||||
"-m",
|
||||
type=str,
|
||||
default="jinavl_reranker",
|
||||
choices=model_example_map.keys(),
|
||||
help="The name of the reranker model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--modality",
|
||||
type=str,
|
||||
default="image",
|
||||
choices=["image", "video"],
|
||||
help="Modality of the multimodal input (image or video).",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def get_multi_modal_input(modality: str) -> tuple[str, ScoreMultiModalParam]:
|
||||
# Sample query for testing the reranker
|
||||
if modality == "image":
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
# Sample multimodal documents to be scored against the query
|
||||
# Each document contains an image URL that will be fetched and processed
|
||||
documents: ScoreMultiModalParam = {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, " # noqa: E501
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust." # noqa: E501
|
||||
),
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
elif modality == "video":
|
||||
query = "A girl is drawing pictures on an ipad."
|
||||
# Sample video documents to be scored against the query
|
||||
documents: ScoreMultiModalParam = {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "A girl is drawing a guitar on her ipad with Apple Pencil.",
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {
|
||||
"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Unsupported modality: {modality}")
|
||||
return query, documents
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
# Run the selected reranker model
|
||||
modality = args.modality
|
||||
model_request = model_example_map[args.model_name](modality)
|
||||
engine_args = model_request.engine_args
|
||||
|
||||
llm = LLM(**asdict(engine_args))
|
||||
|
||||
query, documents = get_multi_modal_input(modality)
|
||||
outputs = llm.score(query, documents, chat_template=model_request.chat_template)
|
||||
|
||||
print("-" * 50)
|
||||
print(f"Model: {engine_args.model}")
|
||||
print(f"Modality: {modality}")
|
||||
print(f"Query: {query}")
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
print("-" * 50)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,71 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example online usage of Score API.
|
||||
|
||||
Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
|
||||
e.g.
|
||||
vllm serve jinaai/jina-reranker-m0 --runner pooling
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-Reranker-2B \
|
||||
--runner pooling \
|
||||
--max-model-len 4096 \
|
||||
--hf_overrides '{"architectures": ["Qwen3VLForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' \
|
||||
--chat-template examples/pooling/score/template/qwen3_vl_reranker.jinja
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
text_1 = "slm markdown"
|
||||
text_2 = {
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
score_url = base_url + "/score"
|
||||
|
||||
response = requests.get(models_url, headers=headers)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
prompt = {"model": model, "text_1": text_1, "text_2": text_2}
|
||||
response = requests.post(score_url, headers=headers, json=prompt)
|
||||
print("\nPrompt when text_1 is string and text_2 is a image list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,54 @@
|
||||
{%- set counter = namespace(index=0) -%}
|
||||
{%- if not tools is defined %}
|
||||
{%- set tools = none %}
|
||||
{%- endif %}
|
||||
|
||||
{%- if messages and messages[0]['role'] == 'system' %}
|
||||
{%- set system_message = messages[0]['content']|trim %}
|
||||
{%- set messages = messages[1:] %}
|
||||
{%- else %}
|
||||
{%- set system_message = "You are a helpful assistant." %}
|
||||
{%- endif %}
|
||||
|
||||
{%- if tools is not none %}
|
||||
{%- set tool_instruction %}
|
||||
You have access to the following tools. When you need to call a tool, you MUST use the following format:
|
||||
|
||||
<tool_call>function_name
|
||||
<arg_key>parameter_name</arg_key>
|
||||
<arg_value>parameter_value</arg_value>
|
||||
</tool_call>
|
||||
|
||||
Important rules:
|
||||
- Always wrap tool calls with <tool_call>...</tool_call> tags
|
||||
- Put the function name on the first line after <tool_call>
|
||||
- Use <arg_key> and <arg_value> tags for each parameter
|
||||
- If a parameter value is a string, keep it as-is. If it's a number or boolean, convert it appropriately
|
||||
- You can make multiple tool calls if needed
|
||||
- If no tool is suitable, respond with regular text
|
||||
|
||||
Available tools:
|
||||
{% endset %}
|
||||
{{- tool_instruction + "\n\n" }}
|
||||
{%- for t in tools %}
|
||||
{{- t | tojson(indent=4) }}
|
||||
{{- "\n\n" }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
|
||||
{%- for message in messages -%}
|
||||
{%- if message['role'] == 'user' -%}
|
||||
{{- '[Round ' + counter.index|string + ']\n问:' + message['content'] -}}
|
||||
{%- set counter.index = counter.index + 1 -%}
|
||||
{%- endif -%}
|
||||
{%- if message['role'] == 'assistant' -%}
|
||||
{{- '\n答:' + message['content'] -}}
|
||||
{%- if (loop.last and add_generation_prompt) or not loop.last -%}
|
||||
{{- '\n' -}}
|
||||
{%- endif -%}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
|
||||
{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
|
||||
{{- '\n答:' -}}
|
||||
{%- endif -%}
|
||||
+3
-1
@@ -80,6 +80,7 @@ plugins:
|
||||
- "re:vllm\\._.*" # Internal modules
|
||||
- "vllm.third_party"
|
||||
- "vllm.vllm_flash_attn"
|
||||
- "re:vllm\\.grpc\\..*_pb2.*" # Auto-generated protobuf files
|
||||
- !ENV [API_AUTONAV_EXCLUDE, "re:^$"] # Match nothing by default
|
||||
- mkdocstrings:
|
||||
handlers:
|
||||
@@ -87,7 +88,8 @@ plugins:
|
||||
options:
|
||||
show_symbol_type_heading: true
|
||||
show_symbol_type_toc: true
|
||||
filters: []
|
||||
filters:
|
||||
- "!.*_pb2_grpc" # Exclude auto-generated gRPC stubs
|
||||
summary:
|
||||
modules: true
|
||||
show_if_no_docstring: true
|
||||
|
||||
@@ -9,6 +9,7 @@ requires = [
|
||||
"torch == 2.9.1",
|
||||
"wheel",
|
||||
"jinja2",
|
||||
"grpcio-tools>=1.76.0",
|
||||
]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
@@ -55,6 +56,10 @@ include = ["vllm*"]
|
||||
"vllm/third_party/**" = ["ALL"]
|
||||
"vllm/version.py" = ["F401"]
|
||||
"vllm/_version.py" = ["ALL"]
|
||||
# Exclude generated protobuf files
|
||||
"vllm/grpc/*_pb2.py" = ["ALL"]
|
||||
"vllm/grpc/*_pb2_grpc.py" = ["ALL"]
|
||||
"vllm/grpc/*_pb2.pyi" = ["ALL"]
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
@@ -162,6 +167,7 @@ depthwise_seperable_CNN = "depthwise_seperable_CNN"
|
||||
[tool.typos.default.extend-words]
|
||||
iy = "iy"
|
||||
tendencias = "tendencias"
|
||||
indx = "indx"
|
||||
# intel cpu features
|
||||
tme = "tme"
|
||||
dout = "dout"
|
||||
|
||||
@@ -9,3 +9,5 @@ wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
build
|
||||
protobuf>=6.33.2
|
||||
grpcio-tools>=1.76.0
|
||||
|
||||
@@ -9,7 +9,7 @@ blake3
|
||||
py-cpuinfo
|
||||
transformers >= 4.56.0, < 5
|
||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||
protobuf # Required by LlamaTokenizer.
|
||||
protobuf >= 6.30.0 # Required by LlamaTokenizer, gRPC.
|
||||
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
|
||||
aiohttp
|
||||
openai >= 1.99.1 # For Responses API with reasoning content
|
||||
@@ -42,7 +42,6 @@ depyf==0.20.0 # required for profiling and debugging with compilation config
|
||||
cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
|
||||
watchfiles # required for http server to monitor the updates of TLS files
|
||||
python-json-logger # Used by logging as per examples/others/logging_configuration.md
|
||||
scipy # Required for phi-4-multimodal-instruct
|
||||
ninja # Required for xgrammar, rocm, tpu, xpu
|
||||
pybase64 # fast base64 implementation
|
||||
cbor2 # Required for cross-language serialization of hashable objects
|
||||
@@ -52,3 +51,5 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
|
||||
anthropic == 0.71.0
|
||||
model-hosting-container-standards >= 0.1.10, < 1.0.0
|
||||
mcp
|
||||
grpcio>=1.76.0
|
||||
grpcio-reflection>=1.76.0
|
||||
@@ -86,3 +86,7 @@ arctic-inference == 0.1.1
|
||||
open-clip-torch==2.32.0
|
||||
# Required for isaac Multi-Modal generation test
|
||||
perceptron==0.1.4
|
||||
# Required for the multi-modal models test
|
||||
timm==1.0.17
|
||||
# Required for plugins test
|
||||
albumentations==1.4.6
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user