forked from Karylab-cklius/vllm
Compare commits
124
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a532c7e0b | ||
|
|
2f3555bf53 | ||
|
|
4699f1bf8b | ||
|
|
681a6371cc | ||
|
|
e86d349053 | ||
|
|
6097afb9bd | ||
|
|
4f4713f96e | ||
|
|
8936118134 | ||
|
|
67ed01c353 | ||
|
|
6e10cb54f6 | ||
|
|
fcb31c1ac3 | ||
|
|
d886c26d4d | ||
|
|
898beca5a8 | ||
|
|
629d45eacb | ||
|
|
f150107efd | ||
|
|
d1135a5087 | ||
|
|
982beae809 | ||
|
|
45232a454e | ||
|
|
03ce1c6ed9 | ||
|
|
4353c9cb4a | ||
|
|
4b7f5ea1a0 | ||
|
|
38907e4391 | ||
|
|
d0359f3e04 | ||
|
|
ed0622e3a8 | ||
|
|
b5f6c5f834 | ||
|
|
bfde49e287 | ||
|
|
153ba7f0f3 | ||
|
|
87518c3027 | ||
|
|
aeee7ef939 | ||
|
|
cda19ecf4d | ||
|
|
80b18230e0 | ||
|
|
d0697cc7b6 | ||
|
|
b0755523dc | ||
|
|
993859ceb0 | ||
|
|
48a65ccb02 | ||
|
|
55842a8d69 | ||
|
|
1f45e83756 | ||
|
|
a8bffaa133 | ||
|
|
5cdddddd4a | ||
|
|
6ef1efd51f | ||
|
|
58da4ee047 | ||
|
|
1ae11e2bfc | ||
|
|
251c18d1f8 | ||
|
|
512765d52d | ||
|
|
640cc9dd7d | ||
|
|
ceade1952c | ||
|
|
747256bb5d | ||
|
|
1174723eba | ||
|
|
6b2b7bd0eb | ||
|
|
70770268c3 | ||
|
|
7a51b3e415 | ||
|
|
d02421a7db | ||
|
|
b1dc87a098 | ||
|
|
79a5b63253 | ||
|
|
c0c98b8b9a | ||
|
|
8d2cff8140 | ||
|
|
4f436782af | ||
|
|
bf45e6d0a5 | ||
|
|
978a4462bb | ||
|
|
1948d0c467 | ||
|
|
4c47710bf7 | ||
|
|
bf9a5ddb24 | ||
|
|
79e799ebbd | ||
|
|
c4e601c73c | ||
|
|
29057d3bee | ||
|
|
219bb5b8c0 | ||
|
|
ad2b1277f9 | ||
|
|
b897f00c9c | ||
|
|
adf9bb3c57 | ||
|
|
b16fda62b7 | ||
|
|
de111f3246 | ||
|
|
afabb5f45a | ||
|
|
3abb7560c0 | ||
|
|
617d1c2ff1 | ||
|
|
692db29cd4 | ||
|
|
82531edbfb | ||
|
|
3daca38e22 | ||
|
|
a302a8fd1b | ||
|
|
4e8c3f1c19 | ||
|
|
5e5afafa21 | ||
|
|
324a3d2bd8 | ||
|
|
4269b79409 | ||
|
|
edc3648966 | ||
|
|
9965f501a8 | ||
|
|
17d87168d2 | ||
|
|
98700c6105 | ||
|
|
10e49d2638 | ||
|
|
8d7c962833 | ||
|
|
f4ddaf8cf7 | ||
|
|
2cdf86044d | ||
|
|
7845379230 | ||
|
|
4b7ca37bd4 | ||
|
|
445b7093fd | ||
|
|
18013df6ae | ||
|
|
c0722f22de | ||
|
|
951dca8019 | ||
|
|
5f7fab881a | ||
|
|
343f65234b | ||
|
|
19fa90ed0d | ||
|
|
03f8d3a548 | ||
|
|
6dc9491406 | ||
|
|
27c0ca50a0 | ||
|
|
7c636432c6 | ||
|
|
c77e596e2e | ||
|
|
ac3dac545b | ||
|
|
39ac640490 | ||
|
|
0b790a2501 | ||
|
|
41488f2acd | ||
|
|
102d51c9f3 | ||
|
|
55e1a8e103 | ||
|
|
21e5a9f48e | ||
|
|
8ad6ff0037 | ||
|
|
f2145efcb6 | ||
|
|
ed33310552 | ||
|
|
3cc328a4be | ||
|
|
3beb57a238 | ||
|
|
8b5531933a | ||
|
|
db8d4a4a06 | ||
|
|
fc701c8058 | ||
|
|
68be0f853e | ||
|
|
60995c05b4 | ||
|
|
29e5d10205 | ||
|
|
235e1f930a | ||
|
+86 |
431cea3eea |
@@ -46,7 +46,7 @@ steps:
|
||||
- tests/models/language/pooling/
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
|
||||
pytest -x -v -s tests/models/language/generation -m cpu_model
|
||||
pytest -x -v -s tests/models/language/pooling -m cpu_model"
|
||||
|
||||
@@ -99,7 +99,7 @@ steps:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
|
||||
pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
|
||||
parallelism: 2
|
||||
parallelism: 3
|
||||
|
||||
- label: "Arm CPU Test"
|
||||
depends_on: []
|
||||
|
||||
@@ -92,8 +92,8 @@ check_and_skip_if_image_exists() {
|
||||
}
|
||||
|
||||
ecr_login() {
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
|
||||
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
|
||||
}
|
||||
|
||||
prepare_cache_tags() {
|
||||
|
||||
@@ -11,7 +11,7 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
|
||||
|
||||
@@ -11,7 +11,7 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
|
||||
|
||||
@@ -98,8 +98,15 @@ steps:
|
||||
commands:
|
||||
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
|
||||
|
||||
- block: "Unblock to build release Docker images"
|
||||
depends_on: ~
|
||||
key: block-build-release-images
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
|
||||
- group: "Build release Docker images"
|
||||
key: "build-release-images"
|
||||
depends_on: block-build-release-images
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
@@ -617,6 +624,8 @@ steps:
|
||||
- label: ":docker: Build release image - x86_64 - ROCm"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: block-build-release-images
|
||||
allow_failure: true
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
|
||||
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
#echo "--- DP+TP"
|
||||
#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
|
||||
#server_pid=$!
|
||||
#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
|
||||
#vllm bench serve \
|
||||
# --backend vllm \
|
||||
# --dataset-name random \
|
||||
# --model meta-llama/Llama-3.2-3B-Instruct \
|
||||
# --num-prompts 20 \
|
||||
# --result-dir ./test_results \
|
||||
# --result-filename dp_pp.json \
|
||||
# --save-result \
|
||||
# --endpoint /v1/completions
|
||||
#kill -s SIGTERM $server_pid; wait $server_pid || true
|
||||
#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
|
||||
#if [ "$failed_req" -ne 0 ]; then
|
||||
# echo "Some requests were failed!"
|
||||
# exit 1
|
||||
#fi
|
||||
echo "--- DP+TP"
|
||||
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
|
||||
server_pid=$!
|
||||
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--dataset-name random \
|
||||
--model meta-llama/Llama-3.2-3B-Instruct \
|
||||
--num-prompts 20 \
|
||||
--result-dir ./test_results \
|
||||
--result-filename dp_pp.json \
|
||||
--save-result \
|
||||
--endpoint /v1/completions
|
||||
kill -s SIGTERM $server_pid; wait $server_pid || true
|
||||
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
|
||||
if [ "$failed_req" -ne 0 ]; then
|
||||
echo "Some requests were failed!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -51,6 +51,7 @@ function cpu_tests() {
|
||||
set -e
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
|
||||
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
|
||||
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
|
||||
|
||||
# basic online serving
|
||||
|
||||
@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
|
||||
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
|
||||
|
||||
# Run the image, setting --shm-size=4g for tensor parallel.
|
||||
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
|
||||
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
|
||||
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
|
||||
|
||||
@@ -769,7 +769,7 @@ steps:
|
||||
- tests/kernels/helion/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pip install helion==0.3.3
|
||||
- pip install helion==1.0.0
|
||||
- pytest -v -s kernels/helion/
|
||||
|
||||
|
||||
@@ -2613,6 +2613,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
|
||||
|
||||
@@ -3601,7 +3602,6 @@ steps:
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
|
||||
|
||||
|
||||
@@ -196,7 +196,8 @@ steps:
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(B200)
|
||||
device: b200
|
||||
|
||||
@@ -141,6 +141,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
|
||||
- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
|
||||
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
|
||||
@@ -155,7 +156,7 @@ steps:
|
||||
- vllm/utils/import_utils.py
|
||||
- tests/kernels/helion/
|
||||
commands:
|
||||
- pip install helion==0.3.3
|
||||
- pip install helion==1.0.0
|
||||
- pytest -v -s kernels/helion/
|
||||
|
||||
|
||||
|
||||
@@ -4,7 +4,6 @@ depends_on:
|
||||
steps:
|
||||
- label: Basic Models Tests (Initialization)
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
torch_nightly: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -73,3 +72,18 @@ steps:
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
commands:
|
||||
- pip install transformers==4.57.5
|
||||
- pytest -v -s tests/models/test_initialization.py
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
|
||||
@@ -42,3 +42,16 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: DFlash Speculators Correctness
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- tests/v1/spec_decode/test_speculators_dflash.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
|
||||
|
||||
+7
-3
@@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
|
||||
/vllm/pooling_params.py @noooop @DarkLight1337
|
||||
/vllm/tokenizers @DarkLight1337 @njhill
|
||||
/vllm/renderers @DarkLight1337 @njhill
|
||||
/vllm/reasoning @aarnphm @chaunceyjiang
|
||||
/vllm/tool_parsers @aarnphm @chaunceyjiang
|
||||
/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
|
||||
/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
|
||||
/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
|
||||
|
||||
# vLLM V1
|
||||
/vllm/v1/attention @LucasWilkinson @MatthewBonanni
|
||||
@@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
|
||||
/tests/v1/kv_connector/nixl_integration @NickLucche
|
||||
/tests/v1/kv_connector @ApostaC @orozery
|
||||
/tests/v1/kv_offload @ApostaC @orozery
|
||||
/tests/v1/determinism @yewentao256
|
||||
/tests/v1/determinism @yewentao256
|
||||
/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
|
||||
/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
|
||||
/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
|
||||
|
||||
# Transformers modeling backend
|
||||
/vllm/model_executor/models/transformers @hmellor
|
||||
|
||||
@@ -15,7 +15,6 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
|
||||
- [ ] The test plan, such as providing test command.
|
||||
- [ ] The test results, such as pasting the results comparison before and after, or e2e results
|
||||
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
|
||||
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0).
|
||||
</details>
|
||||
|
||||
**BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)
|
||||
|
||||
@@ -45,6 +45,7 @@ jobs:
|
||||
- name: Smoke test vllm serve
|
||||
run: |
|
||||
# Start server in background
|
||||
VLLM_CPU_KVCACHE_SPACE=1 \
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--max-model-len=2K \
|
||||
--load-format=dummy \
|
||||
|
||||
@@ -62,14 +62,14 @@ jobs:
|
||||
const prAuthor = context.payload.pull_request.user.login;
|
||||
|
||||
const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
|
||||
q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
|
||||
q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
|
||||
per_page: 1,
|
||||
});
|
||||
|
||||
const authorPRCount = searchResults.total_count;
|
||||
console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
|
||||
const mergedPRCount = searchResults.total_count;
|
||||
console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
|
||||
|
||||
if (authorPRCount === 1) {
|
||||
if (mergedPRCount === 0) {
|
||||
console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
|
||||
await github.rest.issues.createComment({
|
||||
owner,
|
||||
@@ -98,5 +98,5 @@ jobs:
|
||||
].join('\n'),
|
||||
});
|
||||
} else {
|
||||
console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
|
||||
console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
|
||||
}
|
||||
|
||||
+17
-1
@@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${FP4_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
|
||||
# nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
|
||||
# from cache_kernels.cu, so it belongs in _C rather than _C_stable.
|
||||
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${NVFP4_KV_SRC}"
|
||||
CUDA_ARCHS "${FP4_ARCHS}")
|
||||
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
|
||||
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
|
||||
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
|
||||
@@ -944,11 +952,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${FP4_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
|
||||
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${NVFP4_KV_SRC}"
|
||||
CUDA_ARCHS "${FP4_ARCHS}")
|
||||
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
|
||||
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
|
||||
|
||||
+16
-14
@@ -30,6 +30,21 @@ else()
|
||||
list(APPEND CXX_COMPILE_FLAGS
|
||||
"-fopenmp"
|
||||
"-DVLLM_CPU_EXTENSION")
|
||||
|
||||
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
|
||||
# and create a local shim dir with it
|
||||
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
|
||||
|
||||
find_library(OPEN_MP
|
||||
NAMES gomp
|
||||
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
|
||||
NO_DEFAULT_PATH
|
||||
REQUIRED
|
||||
)
|
||||
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
|
||||
if (OPEN_MP)
|
||||
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (NOT MACOSX_FOUND)
|
||||
@@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
|
||||
if(NOT NPROC)
|
||||
set(NPROC 4)
|
||||
endif()
|
||||
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
|
||||
# and create a local shim dir with it
|
||||
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
|
||||
|
||||
find_library(OPEN_MP
|
||||
NAMES gomp
|
||||
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
|
||||
NO_DEFAULT_PATH
|
||||
REQUIRED
|
||||
)
|
||||
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
|
||||
if (OPEN_MP)
|
||||
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
|
||||
endif()
|
||||
|
||||
# Fetch and populate ACL
|
||||
if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
|
||||
@@ -360,6 +361,7 @@ set(VLLM_EXT_SRC
|
||||
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/shm.cpp"
|
||||
"csrc/cpu/activation_lut_bf16.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
|
||||
+22
-2
@@ -724,6 +724,28 @@ void reshape_and_cache_flash(
|
||||
int num_tokens = slot_mapping.size(0);
|
||||
int num_heads = key.size(1);
|
||||
int head_size = key.size(2);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
if (kv_cache_dtype == "nvfp4") {
|
||||
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
|
||||
// NVFP4 dispatch is compiled separately for SM100+.
|
||||
extern void reshape_and_cache_nvfp4_dispatch(
|
||||
torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
|
||||
torch::Tensor & value_cache, torch::Tensor & slot_mapping,
|
||||
torch::Tensor & k_scale, torch::Tensor & v_scale);
|
||||
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
|
||||
slot_mapping, k_scale, v_scale);
|
||||
return;
|
||||
#else
|
||||
TORCH_CHECK(false,
|
||||
"NVFP4 KV cache requires SM100+ (Blackwell). "
|
||||
"Please rebuild vllm with a Blackwell-compatible CUDA target.");
|
||||
#endif
|
||||
}
|
||||
|
||||
// Original FP8/auto path.
|
||||
int block_size = key_cache.size(1);
|
||||
|
||||
int64_t key_stride = key.stride(0);
|
||||
@@ -741,8 +763,6 @@ void reshape_and_cache_flash(
|
||||
|
||||
dim3 grid(num_tokens);
|
||||
dim3 block(std::min(num_heads * head_size, 512));
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
|
||||
CALL_RESHAPE_AND_CACHE_FLASH);
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
|
||||
#include <ATen/ops/empty.h>
|
||||
#include <ATen/ops/gelu.h>
|
||||
#include <c10/util/BFloat16.h>
|
||||
|
||||
constexpr uint32_t ActivationLutSize = 1u << 16;
|
||||
|
||||
at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
|
||||
|
||||
void maybe_init_activation_lut_bf16(
|
||||
uint16_t* lut, std::once_flag& once,
|
||||
at::Tensor (*activation)(const at::Tensor&)) {
|
||||
std::call_once(once, [&]() {
|
||||
auto lut_input =
|
||||
at::empty({static_cast<int64_t>(ActivationLutSize)},
|
||||
at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
|
||||
auto* lut_input_ptr = lut_input.data_ptr<float>();
|
||||
#pragma omp parallel for
|
||||
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
|
||||
lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
|
||||
}
|
||||
|
||||
auto lut_output = activation(lut_input);
|
||||
const auto* lut_output_ptr = lut_output.data_ptr<float>();
|
||||
#pragma omp parallel for
|
||||
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
|
||||
lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
|
||||
const uint16_t* lut, const char* op_name) {
|
||||
TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
|
||||
": input must be bfloat16");
|
||||
TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
|
||||
": out must be bfloat16");
|
||||
TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
|
||||
TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
|
||||
|
||||
const auto* src =
|
||||
reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
|
||||
auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
|
||||
const int64_t n = input.numel();
|
||||
|
||||
CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
dst[i] = lut[src[i]];
|
||||
}
|
||||
CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
|
||||
}
|
||||
|
||||
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
|
||||
const std::string& activation) {
|
||||
if (activation == "gelu") {
|
||||
static std::array<uint16_t, ActivationLutSize> lut{};
|
||||
static std::once_flag once;
|
||||
maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
|
||||
activation_lut_bf16(out, input, lut.data(), "gelu_lut");
|
||||
return;
|
||||
}
|
||||
|
||||
TORCH_CHECK(false, "Unsupported activation: ", activation);
|
||||
}
|
||||
@@ -147,6 +147,9 @@ struct AttentionMetadata {
|
||||
case ISA::NEON:
|
||||
ss << "NEON, ";
|
||||
break;
|
||||
case ISA::VXE:
|
||||
ss << "VXE, ";
|
||||
break;
|
||||
}
|
||||
ss << "workitem_group_num: " << workitem_group_num
|
||||
<< ", reduction_item_num: " << reduction_item_num
|
||||
|
||||
@@ -85,6 +85,9 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
|
||||
at::Tensor& w_scales,
|
||||
std::optional<at::Tensor> bias);
|
||||
|
||||
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
|
||||
const std::string& activation);
|
||||
|
||||
torch::Tensor get_scheduler_metadata(
|
||||
const int64_t num_req, const int64_t num_heads_q,
|
||||
const int64_t num_heads_kv, const int64_t head_dim,
|
||||
@@ -138,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size);
|
||||
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids);
|
||||
|
||||
namespace cpu_utils {
|
||||
void eagle_prepare_inputs_padded_kernel_impl(
|
||||
const torch::Tensor& cu_num_draft_tokens,
|
||||
@@ -231,6 +236,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
|
||||
|
||||
#if (defined(__aarch64__) && !defined(__APPLE__))
|
||||
|
||||
ops.def(
|
||||
"activation_lut_bf16(Tensor! out, Tensor input, str activation)"
|
||||
" -> ()");
|
||||
ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
|
||||
|
||||
#endif // (defined(__aarch64__) && !defined(__APPLE__))
|
||||
|
||||
// Layernorm
|
||||
// Apply Root Mean Square (RMS) Normalization to the input tensor.
|
||||
ops.def(
|
||||
@@ -419,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"block_size) -> ()",
|
||||
&compute_slot_mapping_kernel_impl);
|
||||
|
||||
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
|
||||
|
||||
// Speculative decoding kernels
|
||||
ops.def(
|
||||
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
|
||||
|
||||
+73
-6
@@ -13,13 +13,80 @@
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
#ifdef VLLM_NUMA_DISABLED
|
||||
std::string init_cpu_threads_env(const std::string& cpu_ids) {
|
||||
return std::string(
|
||||
"Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
|
||||
"no effect to setup thread affinity.");
|
||||
}
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
|
||||
#else
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids) {
|
||||
// Memory node binding
|
||||
if (numa_available() != -1) {
|
||||
// Concatenate all node_ids into a single comma-separated string
|
||||
if (!node_ids.empty()) {
|
||||
std::string node_ids_str;
|
||||
for (const int node_id : node_ids) {
|
||||
if (!node_ids_str.empty()) {
|
||||
node_ids_str += ",";
|
||||
}
|
||||
node_ids_str += std::to_string(node_id);
|
||||
}
|
||||
|
||||
#endif
|
||||
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
|
||||
bitmask* src_mask = numa_get_mems_allowed();
|
||||
|
||||
int pid = getpid();
|
||||
|
||||
if (mask && src_mask) {
|
||||
// move all existing pages to the specified numa node.
|
||||
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
|
||||
int page_num = numa_migrate_pages(pid, src_mask, mask);
|
||||
if (page_num == -1) {
|
||||
TORCH_WARN("numa_migrate_pages failed. errno: " +
|
||||
std::to_string(errno));
|
||||
}
|
||||
|
||||
// Restrict memory allocation to the selected NUMA node(s).
|
||||
// Enhances memory locality for the threads bound to those NUMA CPUs.
|
||||
if (node_ids.size() > 1) {
|
||||
errno = 0;
|
||||
numa_set_interleave_mask(mask);
|
||||
if (errno != 0) {
|
||||
TORCH_WARN("numa_set_interleave_mask failed. errno: " +
|
||||
std::to_string(errno));
|
||||
} else {
|
||||
TORCH_WARN(
|
||||
"NUMA binding: Using INTERLEAVE policy for memory "
|
||||
"allocation across multiple NUMA nodes (nodes: " +
|
||||
node_ids_str +
|
||||
"). Memory allocations will be "
|
||||
"interleaved across the specified NUMA nodes.");
|
||||
}
|
||||
} else {
|
||||
errno = 0;
|
||||
numa_set_membind(mask);
|
||||
if (errno != 0) {
|
||||
TORCH_WARN("numa_set_membind failed. errno: " +
|
||||
std::to_string(errno));
|
||||
} else {
|
||||
TORCH_WARN(
|
||||
"NUMA binding: Using MEMBIND policy for memory "
|
||||
"allocation on the NUMA nodes (" +
|
||||
node_ids_str +
|
||||
"). Memory allocations will be "
|
||||
"strictly bound to these NUMA nodes.");
|
||||
}
|
||||
}
|
||||
|
||||
numa_set_strict(1);
|
||||
|
||||
numa_free_nodemask(mask);
|
||||
numa_free_nodemask(src_mask);
|
||||
} else {
|
||||
TORCH_WARN(
|
||||
"numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
|
||||
std::to_string(errno));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif // VLLM_NUMA_DISABLED
|
||||
|
||||
namespace cpu_utils {
|
||||
ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
|
||||
|
||||
@@ -54,12 +54,34 @@ struct Counter {
|
||||
};
|
||||
|
||||
inline int64_t get_available_l2_size() {
|
||||
#if defined(__s390x__)
|
||||
static int64_t size = []() {
|
||||
uint32_t l2_cache_size = 0;
|
||||
auto caps = at::cpu::get_cpu_capabilities();
|
||||
auto it = caps.find("l2_cache_size");
|
||||
if (it != caps.end()) {
|
||||
l2_cache_size = static_cast<uint32_t>(it->second.toInt());
|
||||
}
|
||||
if (l2_cache_size == 0) {
|
||||
long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
|
||||
if (sys_l2 > 0) {
|
||||
l2_cache_size = static_cast<uint32_t>(sys_l2);
|
||||
}
|
||||
}
|
||||
if (l2_cache_size == 0) {
|
||||
l2_cache_size = 256 * 1024;
|
||||
}
|
||||
return static_cast<int64_t>(l2_cache_size) >> 1; // use 50% of L2 cache
|
||||
}();
|
||||
return size;
|
||||
#else
|
||||
static int64_t size = []() {
|
||||
auto caps = at::cpu::get_cpu_capabilities();
|
||||
const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
|
||||
return l2_cache_size >> 1; // use 50% of L2 cache
|
||||
}();
|
||||
return size;
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int32_t alignment_v, typename T>
|
||||
|
||||
@@ -134,4 +134,13 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& input_global_scale);
|
||||
|
||||
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
|
||||
const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b,
|
||||
const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets);
|
||||
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,468 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
|
||||
* Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
|
||||
*/
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include <cutlass/arch/arch.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/epilogue/collective/default_epilogue.hpp"
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/group_array_problem_shape.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
#include <cassert>
|
||||
|
||||
using namespace cute;
|
||||
|
||||
// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
|
||||
template <typename ElementAB, typename ElementC, typename ElementSF,
|
||||
typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
|
||||
__global__ void __mxfp4_get_group_gemm_starts(
|
||||
ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
|
||||
ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
|
||||
LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
|
||||
ElementAB* a_base_as_int, ElementAB* b_base_as_int,
|
||||
ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
|
||||
ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
|
||||
const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
|
||||
int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
|
||||
const int64_t a_stride_val, const int64_t b_stride_val,
|
||||
const int64_t c_stride_val, const int K, const int N) {
|
||||
int64_t expert_id = threadIdx.x;
|
||||
if (expert_id >= gridDim.x * blockDim.x) {
|
||||
return;
|
||||
}
|
||||
int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
|
||||
int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
|
||||
int64_t group_size = 32;
|
||||
int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
|
||||
int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
|
||||
int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
|
||||
assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
|
||||
"unexpected problem sizes");
|
||||
|
||||
int64_t half_k = static_cast<int64_t>(k / 2);
|
||||
int64_t group_k = static_cast<int64_t>(k / group_size);
|
||||
// Shape of A as uint8/byte = [M, K // 2]
|
||||
a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
|
||||
// Shape of B as uint8/byte = [E, N, K // 2]
|
||||
b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
|
||||
// Shape of C = [M, N]
|
||||
out_offsets[expert_id] = out_base_as_int + expert_offset * n;
|
||||
// Shape of a_scale = [sum(sf_sizes), K // group_size]
|
||||
a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
|
||||
|
||||
assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
|
||||
0 &&
|
||||
"TMA requires 128-byte alignment");
|
||||
|
||||
// Shape of B scale = [E, N, K // group_size]
|
||||
b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
|
||||
assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
|
||||
0 &&
|
||||
"TMA requires 128-byte alignment");
|
||||
|
||||
// Initialize strides
|
||||
a_strides[expert_id] = a_stride_val;
|
||||
b_strides[expert_id] = b_stride_val;
|
||||
c_strides[expert_id] = c_stride_val;
|
||||
|
||||
LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
|
||||
LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
|
||||
|
||||
*layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
|
||||
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
|
||||
*layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
|
||||
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
|
||||
}
|
||||
|
||||
#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \
|
||||
TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
|
||||
LayoutSFB, ScaleConfig) \
|
||||
else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
|
||||
__mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
|
||||
LayoutSFB, ScaleConfig> \
|
||||
<<<1, num_experts, 0, stream>>>( \
|
||||
static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
|
||||
static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
|
||||
static_cast<C_TYPE**>(out_starts.data_ptr()), \
|
||||
static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
|
||||
static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
|
||||
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
|
||||
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
|
||||
static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
|
||||
static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
|
||||
static_cast<C_TYPE*>(out_tensors.data_ptr()), \
|
||||
static_cast<SF_TYPE*>(a_scales.data_ptr()), \
|
||||
static_cast<SF_TYPE*>(b_scales.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()), \
|
||||
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>
|
||||
void mxfp4_run_get_group_gemm_starts(
|
||||
const torch::stable::Tensor& a_starts,
|
||||
const torch::stable::Tensor& b_starts,
|
||||
const torch::stable::Tensor& out_starts,
|
||||
const torch::stable::Tensor& a_scales_starts,
|
||||
const torch::stable::Tensor& b_scales_starts,
|
||||
const torch::stable::Tensor& layout_sfa,
|
||||
const torch::stable::Tensor& layout_sfb,
|
||||
const torch::stable::Tensor& a_strides,
|
||||
const torch::stable::Tensor& b_strides,
|
||||
const torch::stable::Tensor& c_strides, int64_t a_stride_val,
|
||||
int64_t b_stride_val, int64_t c_stride_val,
|
||||
torch::stable::Tensor const& a_tensors,
|
||||
torch::stable::Tensor const& b_tensors,
|
||||
torch::stable::Tensor const& out_tensors,
|
||||
torch::stable::Tensor const& a_scales,
|
||||
torch::stable::Tensor const& b_scales,
|
||||
torch::stable::Tensor const& expert_offsets,
|
||||
torch::stable::Tensor const& sf_offsets,
|
||||
torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
|
||||
int num_experts = (int)expert_offsets.size(0);
|
||||
auto stream = get_current_cuda_stream(a_tensors.get_device_index());
|
||||
|
||||
STD_TORCH_CHECK(out_tensors.size(1) == N,
|
||||
"Output tensor shape doesn't match expected shape");
|
||||
STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
|
||||
"b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
|
||||
" dimension must match");
|
||||
if (false) {
|
||||
}
|
||||
// MXFP4 uses E8M0 (float_ue8m0_t) scale factors
|
||||
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
|
||||
torch::headeronly::ScalarType::BFloat16,
|
||||
cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
|
||||
ScaleConfig)
|
||||
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
|
||||
torch::headeronly::ScalarType::Half, half,
|
||||
LayoutSFA, LayoutSFB, ScaleConfig)
|
||||
else {
|
||||
STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutType>
|
||||
void run_mxfp4_blockwise_scaled_group_mm_sm100(
|
||||
torch::stable::Tensor& output, const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
|
||||
using ProblemShape =
|
||||
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
|
||||
using ElementType = cutlass::float_e2m1_t;
|
||||
using ElementSFType = cutlass::float_ue8m0_t;
|
||||
using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
|
||||
using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
|
||||
|
||||
using ElementC = OutType;
|
||||
using ElementD = ElementC;
|
||||
using ElementAccumulator = float;
|
||||
// Layout definitions
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
using LayoutD = LayoutC;
|
||||
|
||||
static constexpr int AlignmentA = 32;
|
||||
static constexpr int AlignmentB = 32;
|
||||
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
|
||||
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
|
||||
|
||||
// Architecture definitions
|
||||
using ArchTag = cutlass::arch::Sm100;
|
||||
using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
|
||||
using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
|
||||
using StageCountType = cutlass::gemm::collective::StageCountAuto;
|
||||
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
struct MMA1SMConfig {
|
||||
using MmaTileShape = Shape<_128, _128, _128>;
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
|
||||
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
|
||||
};
|
||||
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
|
||||
ClusterShape, Shape<_128, _64>, ElementAccumulator,
|
||||
ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
|
||||
LayoutC*, AlignmentD,
|
||||
typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
|
||||
ElementB, LayoutB*, AlignmentB, ElementAccumulator,
|
||||
typename MMA1SMConfig::MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
|
||||
|
||||
using GemmKernel =
|
||||
cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
|
||||
CollectiveEpilogue>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
|
||||
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
|
||||
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
|
||||
using StrideD = typename Gemm::GemmKernel::InternalStrideD;
|
||||
|
||||
using LayoutSFA =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
|
||||
using LayoutSFB =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
|
||||
using ScaleConfig =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
|
||||
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
|
||||
int num_experts = static_cast<int>(expert_offsets.size(0));
|
||||
|
||||
torch::stable::Tensor a_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor b_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor out_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor a_scales_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor b_scales_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor layout_sfa = torch::stable::empty(
|
||||
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
|
||||
a.device());
|
||||
torch::stable::Tensor layout_sfb = torch::stable::empty(
|
||||
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
|
||||
a.device());
|
||||
torch::stable::Tensor a_strides1 =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor b_strides1 =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor c_strides1 =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
|
||||
mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
|
||||
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
|
||||
layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
|
||||
b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
|
||||
b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
|
||||
|
||||
// Create an instance of the GEMM
|
||||
Gemm gemm_op;
|
||||
|
||||
UnderlyingProblemShape* problem_sizes_as_shapes =
|
||||
static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
|
||||
|
||||
// Set the Scheduler info
|
||||
cutlass::KernelHardwareInfo hw_info;
|
||||
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
|
||||
PersistentTileSchedulerSm100GroupParams<
|
||||
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
|
||||
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
|
||||
scheduler.raster_order = RasterOrderOptions::AlongM;
|
||||
hw_info.device_id = a.get_device_index();
|
||||
static std::unordered_map<int, int> cached_sm_counts;
|
||||
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
|
||||
cached_sm_counts[hw_info.device_id] =
|
||||
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
|
||||
hw_info.device_id);
|
||||
}
|
||||
hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
|
||||
|
||||
// Mainloop Arguments
|
||||
typename GemmKernel::MainloopArguments mainloop_args{
|
||||
static_cast<const ElementType**>(a_ptrs.data_ptr()),
|
||||
static_cast<StrideA*>(a_strides1.data_ptr()),
|
||||
static_cast<const ElementType**>(b_ptrs.data_ptr()),
|
||||
static_cast<StrideB*>(b_strides1.data_ptr()),
|
||||
static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
|
||||
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
|
||||
static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
|
||||
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
|
||||
|
||||
// Epilogue Arguments
|
||||
typename GemmKernel::EpilogueArguments epilogue_args{
|
||||
{}, // epilogue.thread
|
||||
nullptr,
|
||||
static_cast<StrideC*>(c_strides1.data_ptr()),
|
||||
static_cast<ElementD**>(out_ptrs.data_ptr()),
|
||||
static_cast<StrideC*>(c_strides1.data_ptr())};
|
||||
auto& fusion_args = epilogue_args.thread;
|
||||
// Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
|
||||
fusion_args.alpha_ptr = nullptr;
|
||||
fusion_args.beta_ptr = nullptr;
|
||||
fusion_args.alpha = 1.0f;
|
||||
fusion_args.alpha_ptr_array = nullptr;
|
||||
fusion_args.dAlpha = {_0{}, _0{}, 0};
|
||||
fusion_args.beta = 0.0f;
|
||||
fusion_args.beta_ptr_array = nullptr;
|
||||
fusion_args.dBeta = {_0{}, _0{}, 0};
|
||||
|
||||
// Gemm Arguments
|
||||
typename GemmKernel::Arguments args{
|
||||
cutlass::gemm::GemmUniversalMode::kGrouped,
|
||||
{num_experts, problem_sizes_as_shapes, nullptr},
|
||||
mainloop_args,
|
||||
epilogue_args,
|
||||
hw_info,
|
||||
scheduler};
|
||||
|
||||
size_t workspace_size = Gemm::get_workspace_size(args);
|
||||
auto workspace =
|
||||
torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
|
||||
std::nullopt, a.device());
|
||||
const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
|
||||
|
||||
auto can_implement_status = gemm_op.can_implement(args);
|
||||
STD_TORCH_CHECK(
|
||||
can_implement_status == cutlass::Status::kSuccess,
|
||||
"Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
|
||||
|
||||
// Run the GEMM
|
||||
auto status = gemm_op.initialize(args, workspace.data_ptr());
|
||||
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
|
||||
"Failed to initialize MXFP4 GEMM: status=", (int)status,
|
||||
" workspace_size=", workspace_size,
|
||||
" num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
|
||||
|
||||
status = gemm_op.run(args, workspace.data_ptr(), stream);
|
||||
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
|
||||
"Failed to run MXFP4 GEMM");
|
||||
}
|
||||
|
||||
template <typename OutType>
|
||||
void run_mxfp4_blockwise_scaled_group_mm(
|
||||
torch::stable::Tensor& output, const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
if (version_num >= 100 && version_num < 120) {
|
||||
run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
|
||||
output, a, b, a_blockscale, b_blockscales, problem_sizes,
|
||||
expert_offsets, sf_offsets, M, N, K);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
STD_TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
|
||||
version_num, ". Required capability: 100");
|
||||
}
|
||||
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
|
||||
// E8M0 scale factors stored as uint8
|
||||
constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
|
||||
#endif
|
||||
|
||||
#define CHECK_TYPE(x, st, m) \
|
||||
STD_TORCH_CHECK(x.scalar_type() == st, \
|
||||
": Inconsistency of torch::stable::Tensor type:", m)
|
||||
#define CHECK_TH_CUDA(x, m) \
|
||||
STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
|
||||
#define CHECK_CONTIGUOUS(x, m) \
|
||||
STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
|
||||
#define CHECK_INPUT(x, st, m) \
|
||||
CHECK_TH_CUDA(x, m); \
|
||||
CHECK_CONTIGUOUS(x, m); \
|
||||
CHECK_TYPE(x, st, m)
|
||||
|
||||
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
|
||||
const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b,
|
||||
const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets) {
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
// Input validation
|
||||
CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
|
||||
CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
|
||||
// MXFP4 uses E8M0 scale factors (stored as uint8)
|
||||
CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
|
||||
CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
|
||||
|
||||
STD_TORCH_CHECK(
|
||||
a_blockscale.dim() == 2,
|
||||
"expected a_blockscale to be of shape [num_experts, rounded_m,"
|
||||
" k // group_size], observed rank: ",
|
||||
a_blockscale.dim())
|
||||
STD_TORCH_CHECK(b_blockscales.dim() == 3,
|
||||
"expected b_blockscale to be of shape: "
|
||||
" [num_experts, n, k // group_size], observed rank: ",
|
||||
b_blockscales.dim())
|
||||
STD_TORCH_CHECK(problem_sizes.dim() == 2,
|
||||
"problem_sizes must be a 2D tensor");
|
||||
STD_TORCH_CHECK(problem_sizes.size(1) == 3,
|
||||
"problem_sizes must have the shape (num_experts, 3)");
|
||||
STD_TORCH_CHECK(
|
||||
problem_sizes.size(0) == expert_offsets.size(0),
|
||||
"Number of experts in problem_sizes must match expert_offsets");
|
||||
STD_TORCH_CHECK(
|
||||
problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"problem_sizes must be int32.");
|
||||
|
||||
int M = static_cast<int>(a.size(0));
|
||||
int N = static_cast<int>(b.size(1));
|
||||
int E = static_cast<int>(b.size(0));
|
||||
int K = static_cast<int>(2 * b.size(2));
|
||||
|
||||
if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
|
||||
output, a, b, a_blockscale, b_blockscales, problem_sizes,
|
||||
expert_offsets, sf_offsets, M, N, K);
|
||||
} else {
|
||||
run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
|
||||
output, a, b, a_blockscale, b_blockscales, problem_sizes,
|
||||
expert_offsets, sf_offsets, M, N, K);
|
||||
}
|
||||
#else
|
||||
STD_TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
|
||||
"SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
|
||||
#endif
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* MXFP4 activation quantization kernel for MoE experts.
|
||||
* Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
|
||||
* over 32-element groups.
|
||||
*
|
||||
* Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
|
||||
* - Block size 32 (2 threads per SF in PACK16 mode)
|
||||
* - E8M0 (power-of-two) scale factors
|
||||
* - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
|
||||
*/
|
||||
|
||||
// MXFP4 requires PACK16 mode (16 elements per thread) so that
|
||||
// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
|
||||
// Must be defined before any header that (transitively) includes
|
||||
// nvfp4_utils.cuh.
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
#include "libtorch_stable/dispatch_utils.h"
|
||||
#include "cuda_vec_utils.cuh"
|
||||
#include "cuda_utils.h"
|
||||
|
||||
#include "nvfp4_utils.cuh"
|
||||
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
|
||||
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
|
||||
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// MXFP4 block size constants
|
||||
static constexpr int MXFP4_SF_VEC_SIZE = 32;
|
||||
|
||||
// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
|
||||
// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
|
||||
static constexpr int MXFP4_NUM_THREADS_PER_SF =
|
||||
MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
// MXFP4 quantization kernel for experts.
|
||||
// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
|
||||
// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
|
||||
// SiLU(gate)*up before quantization.
|
||||
template <class Type, bool FUSE_SILU_MUL = false,
|
||||
bool SMALL_NUM_EXPERTS = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
fp4_packed_t* out, uint32_t* SFout,
|
||||
uint32_t* input_offset_by_experts,
|
||||
uint32_t* output_scale_offset_by_experts,
|
||||
int n_experts, bool low_latency) {
|
||||
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
// MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
|
||||
int32_t const numKTiles = (numCols + 127) / 128;
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
int rowIdx_in_expert = 0;
|
||||
int expert_idx = 0;
|
||||
|
||||
if constexpr (SMALL_NUM_EXPERTS) {
|
||||
for (int i = 0; i < n_experts; i++) {
|
||||
uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
|
||||
uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
|
||||
if (rowIdx >= current_offset && rowIdx < next_offset) {
|
||||
rowIdx_in_expert = rowIdx - current_offset;
|
||||
expert_idx = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
uint32_t local_offsets[17];
|
||||
for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
|
||||
*reinterpret_cast<int4*>(local_offsets) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start]));
|
||||
*reinterpret_cast<int4*>(local_offsets + 4) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start + 4]));
|
||||
*reinterpret_cast<int4*>(local_offsets + 8) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start + 8]));
|
||||
*reinterpret_cast<int4*>(local_offsets + 12) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start + 12]));
|
||||
local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; i++) {
|
||||
if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
|
||||
rowIdx_in_expert = rowIdx - local_offsets[i];
|
||||
expert_idx = chunk_start + i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Load input and optionally apply fused SiLU+Mul
|
||||
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<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;
|
||||
}
|
||||
|
||||
// In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
|
||||
int64_t outOffset = rowIdx * colsPerRow + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
uint32_t* SFout_in_expert =
|
||||
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
|
||||
|
||||
// Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
// Block E8M0 scales only; no extra tensor-level scale in this path
|
||||
constexpr float SFScaleVal = 1.0f;
|
||||
// UE8M0_SF=true for MXFP4 E8M0 scale factors
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
// Large M_topk variant using shared memory for expert offsets
|
||||
template <class Type, bool FUSE_SILU_MUL = false,
|
||||
bool SMALL_NUM_EXPERTS = false>
|
||||
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
fp4_packed_t* out, uint32_t* SFout,
|
||||
uint32_t* input_offset_by_experts,
|
||||
uint32_t* output_scale_offset_by_experts,
|
||||
int n_experts) {
|
||||
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
// MXFP4: numKTiles = ceil(numCols / 128)
|
||||
int32_t const numKTiles = (numCols + 127) / 128;
|
||||
|
||||
extern __shared__ uint32_t shared_input_offsets[];
|
||||
|
||||
if constexpr (SMALL_NUM_EXPERTS) {
|
||||
for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
|
||||
shared_input_offsets[i] = input_offset_by_experts[i];
|
||||
}
|
||||
} else {
|
||||
for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
|
||||
*reinterpret_cast<int4*>(&shared_input_offsets[i]) =
|
||||
*reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
int rowIdx_in_expert = 0;
|
||||
int expert_idx = 0;
|
||||
|
||||
// Binary search through experts using shared memory
|
||||
int left = 0, right = n_experts - 1;
|
||||
while (left <= right) {
|
||||
int mid = (left + right) / 2;
|
||||
uint32_t mid_offset = shared_input_offsets[mid];
|
||||
uint32_t next_offset = shared_input_offsets[mid + 1];
|
||||
|
||||
if (rowIdx >= mid_offset && rowIdx < next_offset) {
|
||||
rowIdx_in_expert = rowIdx - mid_offset;
|
||||
expert_idx = mid;
|
||||
break;
|
||||
} else if (rowIdx < mid_offset) {
|
||||
right = mid - 1;
|
||||
} else {
|
||||
left = mid + 1;
|
||||
}
|
||||
}
|
||||
|
||||
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<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];
|
||||
|
||||
// MXFP4 has no global scale - only block-level E8M0 scale factors
|
||||
constexpr float SFScaleVal = 1.0f;
|
||||
|
||||
uint32_t* SFout_in_expert =
|
||||
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
|
||||
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, bool FUSE_SILU_MUL = false>
|
||||
void mxfp4_quant_impl(void* output, void* output_scale, void* input,
|
||||
void* input_offset_by_experts,
|
||||
void* output_scale_offset_by_experts, int m_topk, int k,
|
||||
int n_experts, cudaStream_t stream) {
|
||||
int multiProcessorCount =
|
||||
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
||||
|
||||
int const workSizePerRow = k / ELTS_PER_THREAD;
|
||||
int const totalWorkSize = m_topk * workSizePerRow;
|
||||
dim3 block(std::min(workSizePerRow, 512));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
|
||||
multiProcessorCount * numBlocksPerSM));
|
||||
while (grid.x <= multiProcessorCount && block.x > 64) {
|
||||
grid.x *= 2;
|
||||
block.x = (block.x + 1) / 2;
|
||||
}
|
||||
|
||||
int const blockRepeat =
|
||||
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
|
||||
if (blockRepeat > 1) {
|
||||
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
|
||||
if (n_experts >= 4) {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_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 {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_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) {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_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 {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
/*Quantization entry for mxfp4 experts quantization*/
|
||||
#define CHECK_TH_CUDA(x, m) \
|
||||
STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x, m) \
|
||||
STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
|
||||
#define CHECK_INPUT(x, m) \
|
||||
CHECK_TH_CUDA(x, m); \
|
||||
CHECK_CONTIGUOUS(x, m);
|
||||
|
||||
constexpr auto HALF = torch::headeronly::ScalarType::Half;
|
||||
constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
|
||||
constexpr auto INT = torch::headeronly::ScalarType::Int;
|
||||
constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
|
||||
|
||||
static constexpr int MXFP4_BLOCK_SIZE = 32;
|
||||
|
||||
static void validate_mxfp4_experts_quant_inputs(
|
||||
torch::stable::Tensor const& output,
|
||||
torch::stable::Tensor const& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts, int64_t m_topk, int64_t k) {
|
||||
CHECK_INPUT(output, "output");
|
||||
CHECK_INPUT(output_scale, "output_scale");
|
||||
CHECK_INPUT(input, "input");
|
||||
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
|
||||
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
|
||||
|
||||
STD_TORCH_CHECK(output.dim() == 2);
|
||||
STD_TORCH_CHECK(output_scale.dim() == 2);
|
||||
STD_TORCH_CHECK(input.dim() == 2);
|
||||
STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
|
||||
STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
|
||||
|
||||
STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
|
||||
STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
|
||||
STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
|
||||
// output is uint8 (two mxfp4 values packed into one uint8)
|
||||
// output_scale is int32 (four E8M0 values packed into one int32)
|
||||
STD_TORCH_CHECK(output.scalar_type() == UINT8);
|
||||
STD_TORCH_CHECK(output_scale.scalar_type() == INT);
|
||||
|
||||
STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
|
||||
STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
|
||||
STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
|
||||
STD_TORCH_CHECK(output.size(0) == m_topk);
|
||||
STD_TORCH_CHECK(output.size(1) == k / 2);
|
||||
int scales_k = k / MXFP4_BLOCK_SIZE;
|
||||
// K-dimension scale columns padded to a multiple of 4 for swizzle layout
|
||||
int padded_k = (scales_k + (4 - 1)) / 4 * 4;
|
||||
// 4 = 4 E8M0 values packed into one int32
|
||||
STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
|
||||
}
|
||||
|
||||
void mxfp4_experts_quant(
|
||||
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts) {
|
||||
auto m_topk = input.size(0);
|
||||
auto k = input.size(1);
|
||||
|
||||
validate_mxfp4_experts_quant_inputs(
|
||||
output, output_scale, input, input_offset_by_experts,
|
||||
output_scale_offset_by_experts, n_experts, m_topk, k);
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
|
||||
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||
input_offset_by_experts.data_ptr(),
|
||||
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
void silu_and_mul_mxfp4_experts_quant(
|
||||
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts) {
|
||||
auto m_topk = input.size(0);
|
||||
auto k_times_2 = input.size(1);
|
||||
STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
|
||||
auto k = k_times_2 / 2;
|
||||
|
||||
validate_mxfp4_experts_quant_inputs(
|
||||
output, output_scale, input, input_offset_by_experts,
|
||||
output_scale_offset_by_experts, n_experts, m_topk, k);
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
|
||||
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||
input_offset_by_experts.data_ptr(),
|
||||
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
|
||||
// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
|
||||
// .cpp files and cannot gate the registration from there.
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
|
||||
m.impl("silu_and_mul_mxfp4_experts_quant",
|
||||
TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
|
||||
}
|
||||
@@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
|
||||
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
|
||||
|
||||
// cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
|
||||
ops.def(
|
||||
"cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
|
||||
" Tensor a_blockscale, Tensor b_blockscales,"
|
||||
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
|
||||
|
||||
// Compute NVFP4 block quantized tensor.
|
||||
ops.def(
|
||||
"scaled_fp4_quant(Tensor input,"
|
||||
@@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts) -> ()");
|
||||
|
||||
// Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
|
||||
ops.def(
|
||||
"mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
|
||||
"Tensor input, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
|
||||
|
||||
// Fused SiLU+Mul+MXFP4 experts quantization.
|
||||
ops.def(
|
||||
"silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
|
||||
"output_scale,"
|
||||
"Tensor input, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
|
||||
|
||||
// Fused SiLU+Mul+NVFP4 quantization.
|
||||
ops.def(
|
||||
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
|
||||
@@ -233,9 +252,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("silu_and_mul_scaled_fp4_experts_quant",
|
||||
TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
|
||||
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
|
||||
|
||||
// W4A8 ops: impl registrations are in the source files
|
||||
// (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
|
||||
// mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
|
||||
// W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -126,7 +126,9 @@ __launch_bounds__(TPB) __global__
|
||||
{
|
||||
const int idx = thread_row_offset + ii;
|
||||
const float val = toFloat(input[idx]);
|
||||
const float softmax_val = expf(val - float_max) * normalizing_factor;
|
||||
float softmax_val = expf(val - float_max) * normalizing_factor;
|
||||
// Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
|
||||
if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
|
||||
output[idx] = softmax_val;
|
||||
}
|
||||
}
|
||||
@@ -147,7 +149,9 @@ __launch_bounds__(TPB) __global__
|
||||
{
|
||||
const int idx = thread_row_offset + ii;
|
||||
const float val = toFloat(input[idx]);
|
||||
const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
|
||||
float sigmoid_val = 1.0f / (1.0f + __expf(-val));
|
||||
// Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
|
||||
if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
|
||||
output[idx] = sigmoid_val;
|
||||
}
|
||||
}
|
||||
@@ -442,6 +446,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
}
|
||||
}
|
||||
|
||||
// Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
|
||||
// NaN gating (from degenerate hidden states in CUDA graph padding) causes
|
||||
// softmax to produce all-NaN, which makes the argmax loop always pick
|
||||
// expert 0 for every top-k slot, producing duplicate expert IDs that
|
||||
// crash FlashInfer's three-step MoE sort.
|
||||
// With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
|
||||
row_chunk[ii] = 0.f;
|
||||
}
|
||||
}
|
||||
|
||||
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
|
||||
|
||||
// If bias is not null, use biased value for selection
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
// NVFP4 KV cache store kernel.
|
||||
// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
|
||||
// into the paged KV cache.
|
||||
//
|
||||
// Per page layout: [K_data | K_scale | V_data | V_scale]
|
||||
// Both data and scale regions are contiguous per head, enabling direct
|
||||
// TMA descriptor use.
|
||||
//
|
||||
// Reuses device functions from nvfp4_utils.cuh:
|
||||
// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale
|
||||
// - pack_fp4() for packing float pairs to fp4
|
||||
// - reciprocal_approximate_ftz() for fast reciprocal
|
||||
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
|
||||
// The swizzle pattern for HND layout is:
|
||||
// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
|
||||
// where T = block_size (page_size), S = scale_dim = head_size // 16.
|
||||
//
|
||||
// For a linear (t, s) position, the swizzled position is:
|
||||
// swizzled_t = (t / 4) * 4 + (s / (S / 4))
|
||||
// swizzled_s = (s % (S / 4)) * 4 + (t % 4)
|
||||
__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
|
||||
int scale_dim) {
|
||||
int s_group = scale_dim / 4;
|
||||
int swizzled_t = (t / 4) * 4 + (s / s_group);
|
||||
int swizzled_s = (s % s_group) * 4 + (t % 4);
|
||||
return swizzled_t * scale_dim + swizzled_s;
|
||||
}
|
||||
|
||||
// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
|
||||
//
|
||||
// Takes separate data and scale cache pointers for K and V.
|
||||
// Within each KV side, data and scale are separate contiguous regions.
|
||||
//
|
||||
// Threading: one CUDA block per token, threads process heads and
|
||||
// groups of 16 elements within each head.
|
||||
template <typename scalar_t>
|
||||
__global__ void reshape_and_cache_nvfp4_kernel(
|
||||
const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
|
||||
const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
|
||||
uint8_t* __restrict__ key_data_cache, // data region for K
|
||||
uint8_t* __restrict__ value_data_cache, // data region for V
|
||||
uint8_t* __restrict__ key_scale_cache, // scale region for K
|
||||
uint8_t* __restrict__ value_scale_cache, // scale region for V
|
||||
const int64_t* __restrict__ slot_mapping, // [num_actual_tokens]
|
||||
const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale
|
||||
const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale
|
||||
const int64_t key_stride, // key.stride(0) in elements
|
||||
const int64_t value_stride, // value.stride(0) in elements
|
||||
const int num_heads, const int head_size, const int block_size,
|
||||
const int64_t data_block_stride, // data cache stride for dim 0
|
||||
const int64_t data_head_stride, // data cache stride for heads
|
||||
const int64_t data_block_offset_stride, // data cache stride for tokens
|
||||
const int64_t scale_block_stride, // scale cache stride for dim 0
|
||||
const int64_t scale_head_stride, // scale cache stride for heads
|
||||
const int64_t scale_block_offset_stride // scale cache stride for tokens
|
||||
) {
|
||||
using CudaType = typename CUDATypeConverter<scalar_t>::Type;
|
||||
using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
|
||||
|
||||
static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8
|
||||
static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
|
||||
|
||||
const int64_t token_idx = blockIdx.x;
|
||||
const int64_t slot_idx = slot_mapping[token_idx];
|
||||
if (slot_idx < 0) return;
|
||||
|
||||
const int64_t block_idx = slot_idx / block_size;
|
||||
const int block_offset = static_cast<int>(slot_idx % block_size);
|
||||
|
||||
const int scale_dim = head_size / 16;
|
||||
const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
|
||||
|
||||
const int total_groups = num_heads * groups_per_head;
|
||||
const int tid = threadIdx.x;
|
||||
const int num_thread_groups = blockDim.x / THREADS_PER_SF;
|
||||
const int tg_id = tid / THREADS_PER_SF;
|
||||
const int tg_lane = tid % THREADS_PER_SF;
|
||||
|
||||
// Process both K (kv=0) and V (kv=1)
|
||||
#pragma unroll
|
||||
for (int kv = 0; kv < 2; kv++) {
|
||||
const scalar_t* __restrict__ src = (kv == 0) ? key : value;
|
||||
const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
|
||||
const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
|
||||
uint8_t* __restrict__ data_cache =
|
||||
(kv == 0) ? key_data_cache : value_data_cache;
|
||||
uint8_t* __restrict__ sc_cache =
|
||||
(kv == 0) ? key_scale_cache : value_scale_cache;
|
||||
|
||||
// Source pointer for this token (use actual stride, not assumed contiguous)
|
||||
const CudaType* __restrict__ token_src =
|
||||
reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
|
||||
|
||||
// Destination bases in data and scale caches for this token's block
|
||||
uint8_t* __restrict__ data_block =
|
||||
data_cache + block_idx * data_block_stride;
|
||||
uint8_t* __restrict__ scale_block =
|
||||
sc_cache + block_idx * scale_block_stride;
|
||||
|
||||
for (int g = tg_id; g < total_groups; g += num_thread_groups) {
|
||||
const int head = g / groups_per_head;
|
||||
const int group_in_head = g % groups_per_head;
|
||||
|
||||
// Load 16 (or 8) bf16 elements from source
|
||||
PVec in_vec;
|
||||
const CudaType* __restrict__ src_ptr =
|
||||
token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
|
||||
tg_lane * ELTS;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < ELTS / 2; i++) {
|
||||
in_vec.elts[i] = reinterpret_cast<
|
||||
const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
|
||||
}
|
||||
|
||||
// Quantize: produces packed fp4 and writes scale factor.
|
||||
uint8_t sf_val;
|
||||
uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
|
||||
|
||||
fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
|
||||
in_vec, global_scale, sf_out_ptr);
|
||||
|
||||
// Write packed FP4 data to data cache
|
||||
uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
|
||||
block_offset * data_block_offset_stride;
|
||||
|
||||
#if CVT_FP4_PACK16
|
||||
{
|
||||
// 16 elements → 8 bytes (u32x2)
|
||||
int data_byte_offset = group_in_head * 8;
|
||||
reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
|
||||
(uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
|
||||
}
|
||||
#else
|
||||
{
|
||||
// 8 elements → 4 bytes (uint32_t)
|
||||
int data_byte_offset =
|
||||
group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
|
||||
reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
|
||||
}
|
||||
#endif
|
||||
|
||||
// Write block scale to scale cache.
|
||||
// K (kv==0): linear layout (no swizzle).
|
||||
// V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
|
||||
if (sf_out_ptr != nullptr) {
|
||||
int scale_idx = group_in_head;
|
||||
uint8_t* __restrict__ scale_dst;
|
||||
if (kv == 0) {
|
||||
scale_dst = scale_block + head * scale_head_stride +
|
||||
block_offset * scale_block_offset_stride + scale_idx;
|
||||
} else {
|
||||
int swizzled_offset =
|
||||
swizzle_scale_offset(block_offset, scale_idx, scale_dim);
|
||||
int swizzled_t = swizzled_offset / scale_dim;
|
||||
int swizzled_s = swizzled_offset % scale_dim;
|
||||
scale_dst = scale_block + head * scale_head_stride +
|
||||
swizzled_t * scale_block_offset_stride + swizzled_s;
|
||||
}
|
||||
*scale_dst = sf_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
// Non-template entry point callable from cache_kernels.cu.
|
||||
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
|
||||
// Each KV side contains both data and scale:
|
||||
// page = [K_data | K_scale | V_data | V_scale]
|
||||
void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
torch::Tensor& slot_mapping,
|
||||
torch::Tensor& k_scale,
|
||||
torch::Tensor& v_scale) {
|
||||
int num_tokens = slot_mapping.size(0);
|
||||
int num_heads = key.size(1);
|
||||
int head_size = key.size(2);
|
||||
int data_dim = head_size / 2;
|
||||
int scale_dim = head_size / 16;
|
||||
int full_dim = data_dim + scale_dim;
|
||||
|
||||
// key_cache is kv_cache[:, 0] with shape
|
||||
// [num_blocks, block_size, num_heads, full_dim] in logical order.
|
||||
// Strides encode the physical layout (HND or NHD).
|
||||
TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
|
||||
TORCH_CHECK(key_cache.size(3) == full_dim,
|
||||
"key_cache last dim must be data_dim + scale_dim, got ",
|
||||
key_cache.size(3), " expected ", full_dim);
|
||||
|
||||
int block_size = key_cache.size(1);
|
||||
|
||||
TORCH_CHECK(head_size % 16 == 0,
|
||||
"head_size must be divisible by 16 for NVFP4 KV cache");
|
||||
TORCH_CHECK(block_size % 4 == 0,
|
||||
"block_size must be divisible by 4 for NVFP4 KV cache swizzle");
|
||||
|
||||
// Detect physical layout from strides (based on full_dim).
|
||||
// HND: head stride > block_offset stride.
|
||||
bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
|
||||
|
||||
int64_t data_block_stride = key_cache.stride(0); // page_bytes
|
||||
int64_t data_head_stride, data_block_offset_stride;
|
||||
if (is_hnd) {
|
||||
data_head_stride = (int64_t)block_size * data_dim;
|
||||
data_block_offset_stride = data_dim;
|
||||
} else {
|
||||
data_head_stride = data_dim;
|
||||
data_block_offset_stride = (int64_t)num_heads * data_dim;
|
||||
}
|
||||
|
||||
// Page layout: [K_data | K_scale | V_data | V_scale]
|
||||
// Scale follows data within each KV side.
|
||||
int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
|
||||
|
||||
uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
|
||||
uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
|
||||
|
||||
// Scale strides: same page stride, inner strides from layout.
|
||||
int64_t scale_block_stride = data_block_stride;
|
||||
int64_t scale_head_stride, scale_block_offset_stride;
|
||||
if (is_hnd) {
|
||||
scale_head_stride = (int64_t)block_size * scale_dim;
|
||||
scale_block_offset_stride = scale_dim;
|
||||
} else {
|
||||
scale_head_stride = scale_dim;
|
||||
scale_block_offset_stride = (int64_t)num_heads * scale_dim;
|
||||
}
|
||||
|
||||
const float* k_scale_ptr = k_scale.data_ptr<float>();
|
||||
const float* v_scale_ptr = v_scale.data_ptr<float>();
|
||||
|
||||
int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
|
||||
int total_groups = num_heads * groups_per_head;
|
||||
constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
|
||||
int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
|
||||
num_threads = ((num_threads + 31) / 32) * 32;
|
||||
|
||||
dim3 grid(num_tokens);
|
||||
dim3 block(num_threads);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(
|
||||
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
|
||||
vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
|
||||
key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
|
||||
key_scale_ptr, value_scale_ptr,
|
||||
slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
|
||||
key.stride(0), value.stride(0), num_heads, head_size,
|
||||
block_size, data_block_stride, data_head_stride,
|
||||
data_block_offset_stride, scale_block_stride, scale_head_stride,
|
||||
scale_block_offset_stride);
|
||||
});
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
#pragma once
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <torch/library.h>
|
||||
#include <tuple>
|
||||
|
||||
|
||||
+5
-4
@@ -642,7 +642,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
else \
|
||||
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
|
||||
fi; \
|
||||
uv pip install --system accelerate hf_transfer modelscope \
|
||||
uv pip install --system accelerate modelscope \
|
||||
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
|
||||
|
||||
# ============================================================
|
||||
@@ -756,9 +756,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -e tests/vllm_test_utils
|
||||
|
||||
# enable fast downloads from hf (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system hf_transfer
|
||||
ENV HF_HUB_ENABLE_HF_TRANSFER 1
|
||||
ENV HF_XET_HIGH_PERFORMANCE 1
|
||||
|
||||
# increase timeout for hf downloads (for testing)
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
||||
|
||||
# Copy in the v1 package for testing (it isn't distributed yet)
|
||||
COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
|
||||
|
||||
@@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/dev.txt && \
|
||||
uv pip install -r requirements/lint.txt && \
|
||||
uv pip install -r requirements/test/cpu.txt && \
|
||||
pre-commit install --hook-type pre-commit --hook-type commit-msg
|
||||
|
||||
ENTRYPOINT ["bash"]
|
||||
@@ -197,6 +198,12 @@ ADD ./.buildkite/ ./.buildkite/
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# enable fast downloads from hf (for testing)
|
||||
ENV HF_XET_HIGH_PERFORMANCE 1
|
||||
|
||||
# increase timeout for hf downloads (for testing)
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
||||
|
||||
######################### RELEASE IMAGE #########################
|
||||
FROM base AS vllm-openai
|
||||
|
||||
|
||||
@@ -272,9 +272,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -e tests/vllm_test_utils
|
||||
|
||||
# enable fast downloads from hf (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system hf_transfer
|
||||
ENV HF_HUB_ENABLE_HF_TRANSFER 1
|
||||
ENV HF_XET_HIGH_PERFORMANCE 1
|
||||
|
||||
# increase timeout for hf downloads (for testing)
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/test/nightly-torch.txt
|
||||
|
||||
+11
-11
@@ -192,9 +192,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
|
||||
FROM base AS build_deep
|
||||
ARG ROCSHMEM_BRANCH="ba0bf0f3"
|
||||
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
|
||||
ARG DEEPEP_BRANCH="e84464ec"
|
||||
ARG DEEPEP_BRANCH="5d90af8b"
|
||||
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
|
||||
ARG DEEPEP_NIC="cx7"
|
||||
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
|
||||
ENV ROCSHMEM_DIR=/opt/rocshmem
|
||||
|
||||
RUN git clone ${ROCSHMEM_REPO} \
|
||||
@@ -202,13 +203,11 @@ RUN git clone ${ROCSHMEM_REPO} \
|
||||
&& git checkout ${ROCSHMEM_BRANCH} \
|
||||
&& mkdir -p projects/rocshmem/build \
|
||||
&& cd projects/rocshmem/build \
|
||||
&& cmake .. \
|
||||
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
|
||||
-DROCM_PATH=/opt/rocm \
|
||||
-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
|
||||
-DUSE_EXTERNAL_MPI=OFF \
|
||||
&& make -j \
|
||||
&& make install
|
||||
&& bash ../scripts/build_configs/all_backends \
|
||||
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
|
||||
-DROCM_PATH=/opt/rocm \
|
||||
-DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
|
||||
-DUSE_EXTERNAL_MPI=OFF
|
||||
|
||||
# Build DeepEP wheel.
|
||||
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
|
||||
@@ -365,9 +364,10 @@ RUN cd /vllm-workspace \
|
||||
&& python3 -m pip install pytest-shard
|
||||
|
||||
# enable fast downloads from hf (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system hf_transfer
|
||||
ENV HF_HUB_ENABLE_HF_TRANSFER=1
|
||||
ENV HF_XET_HIGH_PERFORMANCE=1
|
||||
|
||||
# increase timeout for hf downloads (for testing)
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
||||
|
||||
# install audio decode package `torchcodec` from source (required due to
|
||||
# ROCm and torch version mismatch) for tests with datasets package
|
||||
|
||||
+34
-34
@@ -42,7 +42,7 @@ FROM python-install AS pyarrow
|
||||
# Build Apache Arrow
|
||||
WORKDIR /tmp
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/apache/arrow.git && \
|
||||
git clone https://github.com/apache/arrow.git -b maint-19.0.1 && \
|
||||
cd arrow/cpp && \
|
||||
mkdir release && cd release && \
|
||||
cmake -DCMAKE_BUILD_TYPE=Release \
|
||||
@@ -68,19 +68,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements-build.txt && \
|
||||
python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
|
||||
|
||||
FROM python-install AS numa-build
|
||||
# Install numactl (needed for numa.h dependency)
|
||||
WORKDIR /tmp
|
||||
RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \
|
||||
tar -xvzf v2.0.16.tar.gz && \
|
||||
cd numactl-2.0.16 && \
|
||||
./autogen.sh && \
|
||||
./configure && \
|
||||
make
|
||||
|
||||
# Set include path
|
||||
ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
|
||||
|
||||
FROM python-install AS rust
|
||||
ENV CARGO_HOME=/root/.cargo
|
||||
ENV RUSTUP_HOME=/root/.rustup
|
||||
@@ -91,6 +78,18 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \
|
||||
rustup default stable && \
|
||||
rustup show
|
||||
|
||||
FROM python-install AS numa-build
|
||||
WORKDIR /tmp
|
||||
RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \
|
||||
tar -xvzf v2.0.19.tar.gz && \
|
||||
cd numactl-2.0.19 && \
|
||||
./autogen.sh && \
|
||||
./configure && \
|
||||
make
|
||||
|
||||
# Set include path
|
||||
ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
|
||||
|
||||
FROM python-install AS torch-vision
|
||||
# Install torchvision
|
||||
ARG TORCH_VISION_VERSION=v0.26.0
|
||||
@@ -133,7 +132,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \
|
||||
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
|
||||
cd llvm-project && mkdir build && cd build && \
|
||||
uv pip install 'cmake<4' setuptools numpy && \
|
||||
uv pip install 'cmake<4' 'setuptools<70' numpy && \
|
||||
export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \
|
||||
CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \
|
||||
CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \
|
||||
@@ -193,27 +192,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
cd opencv-python && \
|
||||
python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist
|
||||
|
||||
# Build Outlines Core
|
||||
FROM python-install AS outlines-core-builder
|
||||
## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330
|
||||
FROM python-install AS guidance-builder
|
||||
WORKDIR /tmp
|
||||
ENV CARGO_HOME=/root/.cargo
|
||||
ENV RUSTUP_HOME=/root/.rustup
|
||||
ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
|
||||
COPY requirements/common.txt /tmp/requirements/common.txt
|
||||
ARG OUTLINES_CORE_VERSION
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
|
||||
--mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
|
||||
OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \
|
||||
if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \
|
||||
git clone https://github.com/dottxt-ai/outlines-core.git && \
|
||||
cd outlines-core && \
|
||||
git checkout tags/${OUTLINES_CORE_VERSION} && \
|
||||
sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \
|
||||
git clone https://github.com/guidance-ai/llguidance.git && \
|
||||
cd llguidance && \
|
||||
git checkout s390x-fix-v2 && \
|
||||
uv pip install maturin && \
|
||||
python -m maturin build --release --out dist
|
||||
python -m maturin build --release --out dist --compatibility linux
|
||||
|
||||
# Final build stage
|
||||
# # Final build stage
|
||||
FROM python-install AS vllm-cpu
|
||||
ARG PYTHON_VERSION
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
|
||||
@@ -229,10 +223,12 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/
|
||||
ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
|
||||
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
# Force pure Python protobuf to avoid s390x C++ extension crashes
|
||||
ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
|
||||
COPY . /workspace/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \
|
||||
RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
|
||||
make -C /numactl install
|
||||
|
||||
# Install dependencies, including PyTorch and Apache Arrow
|
||||
@@ -245,22 +241,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \
|
||||
--mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \
|
||||
--mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \
|
||||
--mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \
|
||||
ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \
|
||||
--mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \
|
||||
ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \
|
||||
VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \
|
||||
HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \
|
||||
LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \
|
||||
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
|
||||
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
|
||||
OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \
|
||||
uv pip install -v \
|
||||
$ARROW_WHL_FILE \
|
||||
GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
|
||||
uv pip install -v \
|
||||
$ARROW_WHL_FILE \
|
||||
$VISION_WHL_FILE \
|
||||
$HF_XET_WHL_FILE \
|
||||
$LLVM_WHL_FILE \
|
||||
$NUMBA_WHL_FILE \
|
||||
$OPENCV_WHL_FILE \
|
||||
$OUTLINES_CORE_WHL_FILE \
|
||||
$GUIDANCE_WHL_FILE \
|
||||
--index-strategy unsafe-best-match \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt
|
||||
@@ -271,6 +267,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
|
||||
uv pip install "$(echo dist/*.whl)[tensorizer]"
|
||||
|
||||
# Remove protobuf C++ extension that crashes on s390x
|
||||
RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \
|
||||
/opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true
|
||||
|
||||
# setup non-root user for vllm
|
||||
RUN umask 002 && \
|
||||
/usr/sbin/useradd --uid 2000 --gid 0 vllm && \
|
||||
|
||||
@@ -37,6 +37,7 @@ th {
|
||||
| HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
|
||||
| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` |
|
||||
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
|
||||
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
|
||||
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
|
||||
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
|
||||
|
||||
@@ -239,6 +240,69 @@ vllm bench serve \
|
||||
--spec-bench-category "summarization"
|
||||
```
|
||||
|
||||
#### SPEED-Bench Benchmark with Speculative Decoding
|
||||
|
||||
[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k).
|
||||
|
||||
!!! note
|
||||
This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets.
|
||||
|
||||
First, download the dataset to a folder, using this one liner:
|
||||
|
||||
```bash
|
||||
curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
|
||||
```
|
||||
|
||||
The command supports also the following arguments:
|
||||
|
||||
- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets.
|
||||
- `--output_dir`: download to a specified folder. By default, it will download to the current directory.
|
||||
|
||||
Start a server with speculative decoding:
|
||||
|
||||
```bash
|
||||
vllm serve meta-llama/Llama-3.3-70B-Instruct \
|
||||
--speculative-config $'{"method": "eagle3",
|
||||
"num_speculative_tokens": 3,
|
||||
"model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}'
|
||||
```
|
||||
|
||||
Run all categories in the Qualitative split:
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model meta-llama/Llama-3.3-70B-Instruct \
|
||||
--dataset-name speed_bench \
|
||||
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
|
||||
--num-prompts -1
|
||||
```
|
||||
|
||||
Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`.
|
||||
|
||||
Run only a specific category like "multilingual":
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model meta-llama/Llama-3.3-70B-Instruct \
|
||||
--dataset-name speed_bench \
|
||||
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
|
||||
--num-prompts -1
|
||||
--speed-bench-category "multilingual"
|
||||
```
|
||||
|
||||
Run all categories in the Throughput split (2k ISL):
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model meta-llama/Llama-3.3-70B-Instruct \
|
||||
--dataset-name speed_bench \
|
||||
--speed-bench-dataset-subset throughput_2k
|
||||
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
|
||||
--num-prompts -1
|
||||
```
|
||||
|
||||
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
|
||||
|
||||
#### Other HuggingFaceDataset Examples
|
||||
|
||||
```bash
|
||||
|
||||
@@ -193,7 +193,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
|
||||
|
||||
The API server takes care of basic audio I/O and optional chunking before building prompts:
|
||||
|
||||
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`.
|
||||
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
|
||||
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
|
||||
- Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
|
||||
|
||||
@@ -206,8 +206,8 @@ Relevant server logic:
|
||||
async def _preprocess_speech_to_text(...):
|
||||
language = self.model_cls.validate_language(request.language)
|
||||
...
|
||||
y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate)
|
||||
duration = librosa.get_duration(y=y, sr=sr)
|
||||
y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate)
|
||||
duration = get_audio_duration(y=y, sr=sr)
|
||||
do_split_audio = (self.asr_config.allow_audio_chunking
|
||||
and duration > self.asr_config.max_audio_clip_s)
|
||||
chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
|
||||
|
||||
@@ -206,8 +206,8 @@ Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_cont
|
||||
used to profile a section of code.
|
||||
|
||||
!!! note
|
||||
The legacy import paths `vllm.utils.cprofile` and `vllm.utils.cprofile_context` are deprecated.
|
||||
Please use `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` instead.
|
||||
The `vllm.utils.profiling` helpers are deprecated and will be removed in
|
||||
`v0.21`. Please use Python's `cProfile` module directly instead.
|
||||
|
||||
### Example usage - decorator
|
||||
|
||||
|
||||
@@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first).
|
||||
| 2 | `FLASH_ATTN` |
|
||||
| 3 | `TRITON_ATTN` |
|
||||
| 4 | `FLEX_ATTENTION` |
|
||||
| 5 | `TURBOQUANT` |
|
||||
|
||||
**Ampere/Hopper (SM 8.x-9.x):**
|
||||
|
||||
@@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first).
|
||||
| 2 | `FLASHINFER` |
|
||||
| 3 | `TRITON_ATTN` |
|
||||
| 4 | `FLEX_ATTENTION` |
|
||||
| 5 | `TURBOQUANT` |
|
||||
|
||||
### MLA Attention (DeepSeek-style)
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ or just on the low or high end.
|
||||
| Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` |
|
||||
| ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
|
||||
| [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |
|
||||
| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low |
|
||||
| [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always |
|
||||
| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always |
|
||||
| [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |
|
||||
@@ -40,6 +41,7 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
| Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm |
|
||||
| ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
|
||||
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
|
||||
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
|
||||
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
|
||||
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
|
||||
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
|
||||
@@ -54,6 +56,9 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
|
||||
for per-backend details.
|
||||
|
||||
\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
|
||||
tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
|
||||
|
||||
† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
|
||||
other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
|
||||
SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
|
||||
@@ -184,6 +189,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
|
||||
|
||||
### MiniMax QK Norm (`fuse_minimax_qk_norm`)
|
||||
|
||||
!!! info
|
||||
This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
|
||||
the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
|
||||
and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
|
||||
optimization level.
|
||||
|
||||
**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
|
||||
per-token Q/K variances before applying RMS normalization to Q and K.
|
||||
|
||||
This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
|
||||
`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
|
||||
`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve MiniMaxAI/MiniMax-M2.5 \
|
||||
--tensor-parallel-size 4 \
|
||||
--compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
|
||||
```
|
||||
|
||||
**Code locations.**
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
|
||||
- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
|
||||
- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
|
||||
|
||||
### Sequence Parallelism (`enable_sp`)
|
||||
|
||||
**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
|
||||
|
||||
@@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us
|
||||
|
||||
### Grafana Dashboard
|
||||
|
||||
vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
|
||||
vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
|
||||
|
||||
The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important:
|
||||
|
||||
@@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing:
|
||||
- Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
|
||||
- Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
|
||||
- [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
|
||||
- [User-facing docs](../../examples/online_serving/opentelemetry/README.md)
|
||||
- [User-facing docs](../../examples/observability/opentelemetry/README.md)
|
||||
- [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
|
||||
- [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview)
|
||||
|
||||
|
||||
@@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl
|
||||
Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech.
|
||||
|
||||
```python
|
||||
import librosa
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.multimodal.audio import split_audio
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
# Load long audio file
|
||||
audio, sr = librosa.load("long_audio.wav", sr=16000)
|
||||
audio, sr = load_audio("long_audio.wav", sr=16000)
|
||||
|
||||
# Split into chunks at low-energy (quiet) regions
|
||||
chunks = split_audio(
|
||||
@@ -832,7 +832,7 @@ Then, you can use the OpenAI client as follows:
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
# Any format supported by librosa is supported
|
||||
# Any format supported by soundfile/PyAV is supported
|
||||
audio_url = AudioAsset("winning_call").url
|
||||
audio_base64 = encode_base64_content_from_url(audio_url)
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ The following are the supported quantization formats for vLLM:
|
||||
- [INT8 W8A8](int8.md)
|
||||
- [FP8 W8A8](fp8.md)
|
||||
- [NVIDIA Model Optimizer](modelopt.md)
|
||||
- [Online Quantization](online.md)
|
||||
- [AMD Quark](quark.md)
|
||||
- [Quantized KV Cache](quantized_kvcache.md)
|
||||
- [TorchAO](torchao.md)
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# Online Quantization
|
||||
|
||||
Online quantization lets you take a BF16/FP16 model and quantize its Linear
|
||||
and MoE weights to lower precision (such as FP8) at load time, without needing
|
||||
a pre-quantized checkpoint or calibration data. Weights are converted during
|
||||
model loading and activations are dynamically scaled during each forward pass.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Pass a scheme name to the `quantization` parameter:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
# Per-tensor FP8 quantization (one scale per weight tensor)
|
||||
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
|
||||
|
||||
# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
|
||||
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
|
||||
```
|
||||
|
||||
Or with the CLI:
|
||||
|
||||
```bash
|
||||
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
|
||||
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
|
||||
```
|
||||
|
||||
## Supported Schemes
|
||||
|
||||
| Scheme | Weight recipe | Activation recipe | Notes |
|
||||
| ------ | ------------- | ------------------ | ----- |
|
||||
| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
|
||||
| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
|
||||
|
||||
Support for additional schemes will be added in future versions of vllm.
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
For fine-grained control, use a `quantization_config` dictionary.
|
||||
|
||||
### Separate Schemes for Dense and MoE Layers
|
||||
|
||||
You can apply different quantization schemes to dense linear layers and MoE expert layers:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM(
|
||||
"ibm-granite/granite-3.0-1b-a400m-base",
|
||||
quantization="fp8_per_tensor",
|
||||
quantization_config={
|
||||
"linear_scheme_override": "fp8_per_block",
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
Or,
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM(
|
||||
"ibm-granite/granite-3.0-1b-a400m-base",
|
||||
quantization="fp8_per_tensor",
|
||||
quantization_config={
|
||||
"moe_scheme_override": "fp8_per_block",
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Excluding Layers from Quantization
|
||||
|
||||
Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`):
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM(
|
||||
"ibm-granite/granite-3.0-1b-a400m-base",
|
||||
quantization="fp8_per_tensor",
|
||||
quantization_config={
|
||||
"ignore": [
|
||||
# exact layer name
|
||||
"model.layers.1.self_attn.o_proj",
|
||||
# regex: skip all QKV projections
|
||||
"re:.*[qkv]_proj",
|
||||
],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
!!! note
|
||||
For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name.
|
||||
@@ -3,15 +3,15 @@
|
||||
|
||||
vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform.
|
||||
|
||||
Currently, the CPU implementation for s390x architecture supports FP32 datatype only.
|
||||
Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16.
|
||||
|
||||
--8<-- [end:installation]
|
||||
--8<-- [start:requirements]
|
||||
|
||||
- OS: `Linux`
|
||||
- SDK: `gcc/g++ >= 12.3.0` or later with Command Line Tools
|
||||
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
|
||||
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
|
||||
- Build install python packages: `pyarrow`, `torch` and `torchvision`
|
||||
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
|
||||
|
||||
--8<-- [end:requirements]
|
||||
--8<-- [start:set-up-using-python]
|
||||
@@ -24,13 +24,14 @@ Currently, there are no pre-built IBM Z CPU wheels.
|
||||
--8<-- [end:pre-built-wheels]
|
||||
--8<-- [start:build-wheel-from-source]
|
||||
|
||||
Install the following packages from the package manager before building the vLLM. For example on RHEL 9.4:
|
||||
Install the following packages from the package manager before building the vLLM. For example on RHEL 9.6:
|
||||
|
||||
```bash
|
||||
dnf install -y \
|
||||
which procps findutils tar vim git gcc g++ make patch make cython zlib-devel \
|
||||
which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
|
||||
libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
|
||||
openssl-devel openblas openblas-devel wget autoconf automake libtool cmake numactl-devel
|
||||
openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
|
||||
clang llvm-devel llvm-static clang-devel
|
||||
```
|
||||
|
||||
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
|
||||
@@ -43,13 +44,13 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
|
||||
Execute the following commands to build and install vLLM from source.
|
||||
|
||||
!!! tip
|
||||
Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM.
|
||||
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
|
||||
|
||||
```bash
|
||||
sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds
|
||||
uv pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
--torch-backend auto \
|
||||
-r requirements/build/cuda.txt \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
uv pip install dist/*.whl
|
||||
@@ -57,10 +58,9 @@ Execute the following commands to build and install vLLM from source.
|
||||
|
||||
??? console "pip"
|
||||
```bash
|
||||
sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds
|
||||
pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/nightly/cpu \
|
||||
-r requirements/build/cuda.txt \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
pip install dist/*.whl
|
||||
|
||||
@@ -240,7 +240,7 @@ uv pip install vllm==${VLLM_VERSION} \
|
||||
# Install dependencies
|
||||
pip install --upgrade numba \
|
||||
scipy \
|
||||
huggingface-hub[cli,hf_transfer] \
|
||||
huggingface-hub[cli] \
|
||||
setuptools_scm
|
||||
pip install -r requirements/rocm.txt
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle:
|
||||
- [@aarnphm](https://github.com/aarnphm): Structured output
|
||||
- [@alexm-redhat](https://github.com/alexm-redhat): Performance
|
||||
- [@ApostaC](https://github.com/ApostaC): Connectors, offloading
|
||||
- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
|
||||
- [@benchislett](https://github.com/benchislett): Engine core and spec decode
|
||||
- [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
|
||||
- [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
|
||||
@@ -31,6 +32,7 @@ Sorted alphabetically by GitHub handle:
|
||||
- [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
|
||||
- [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed
|
||||
- [@markmc](https://github.com/markmc): Observability
|
||||
- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance
|
||||
- [@mgoin](https://github.com/mgoin): Quantization and performance
|
||||
- [@NickLucche](https://github.com/NickLucche): KV connector
|
||||
- [@njhill](https://github.com/njhill): Distributed, API server, engine core
|
||||
@@ -41,6 +43,7 @@ Sorted alphabetically by GitHub handle:
|
||||
- [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
|
||||
- [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
|
||||
- [@russellb](https://github.com/russellb): Structured output, engine core, security
|
||||
- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
|
||||
- [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
|
||||
- [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
|
||||
- [@tdoublep](https://github.com/tdoublep): State space models
|
||||
@@ -86,7 +89,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
|
||||
- AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
|
||||
- @robertgshaw2-redhat, @njhill, @russellb
|
||||
- ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
|
||||
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
|
||||
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
|
||||
- KV Connector: Connector interface and implementation for KV cache offload and transfer
|
||||
- @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC
|
||||
- Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks
|
||||
@@ -105,7 +108,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
|
||||
- Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms
|
||||
- @ProExpertProg
|
||||
- Attention: Attention interface for paged attention
|
||||
- @WoosukKwon, @LucasWilkinson, @heheda12345
|
||||
- @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni
|
||||
- FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
|
||||
- @tlrmchlsmth
|
||||
- Quantization: Various quantization config, weight loading, and kernel.
|
||||
@@ -119,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
|
||||
- State space models: The state space models implementation in vLLM
|
||||
- @tdoublep, @tlrmchlsmth
|
||||
- Reasoning and tool calling parsers
|
||||
- @chaunceyjiang, @aarnphm
|
||||
- @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning
|
||||
|
||||
### Entrypoints
|
||||
|
||||
@@ -133,7 +136,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
|
||||
### Features
|
||||
|
||||
- Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
|
||||
- @WoosukKwon, @benchislett, @luccafong
|
||||
- @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni
|
||||
- Structured Output: The structured output implementation
|
||||
- @russellb, @aarnphm
|
||||
- RL: The RL related features such as collective rpc, sleep mode, etc.
|
||||
@@ -153,8 +156,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
|
||||
|
||||
### External Kernels Integration
|
||||
|
||||
- FlashAttention: @LucasWilkinson
|
||||
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
|
||||
- FlashAttention: @LucasWilkinson, @MatthewBonanni
|
||||
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
|
||||
- Blackwell Kernels: @mgoin, @yewentao256
|
||||
- DeepEP/DeepGEMM: @mgoin, @yewentao256
|
||||
|
||||
|
||||
@@ -59,6 +59,16 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
|
||||
Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
|
||||
are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
|
||||
|
||||
### Pooling Types
|
||||
|
||||
| Pooling Tasks | Granularity | Description |
|
||||
|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `CLS` pooling | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. |
|
||||
| `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output. |
|
||||
| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling. |
|
||||
| `ALL` pooling | Token-wise | Outputs the last_hidden_states for all input tokens. |
|
||||
| `STEP` pooling | Token-wise | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids. |
|
||||
|
||||
### Score Types
|
||||
|
||||
The scoring models is designed to compute similarity scores between two input prompts. It supports three model types
|
||||
|
||||
@@ -45,6 +45,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th
|
||||
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
|
||||
| `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | |
|
||||
| `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | |
|
||||
| `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ |
|
||||
| `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
|
||||
| `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
|
||||
@@ -73,6 +74,12 @@ You can compute pairwise similarity scores to build a similarity matrix using th
|
||||
!!! note
|
||||
`jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights.
|
||||
|
||||
!!! note
|
||||
`jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters
|
||||
(`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the
|
||||
selected adapter into the base weights at load time. Choose the task with
|
||||
`--hf-overrides '{"jina_task": "<task>"}'`; the default is `retrieval`.
|
||||
|
||||
### Multimodal Models
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -160,6 +160,8 @@ The following Score API parameters are supported:
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
|
||||
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
|
||||
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params"
|
||||
```
|
||||
|
||||
#### Examples
|
||||
@@ -370,6 +372,8 @@ The following rerank api parameters are supported:
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
|
||||
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
|
||||
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params"
|
||||
```
|
||||
|
||||
#### Examples
|
||||
|
||||
@@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the
|
||||
Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
|
||||
Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
|
||||
|
||||
### As Reward Models
|
||||
### Reward Models
|
||||
|
||||
Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
|
||||
|
||||
|
||||
@@ -9,14 +9,14 @@
|
||||
- Online APIs:
|
||||
- Pooling API (`/pooling`)
|
||||
|
||||
The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs a embedding for each token.
|
||||
The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token.
|
||||
|
||||
Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md).
|
||||
|
||||
!!! note
|
||||
|
||||
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not
|
||||
what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or
|
||||
what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or
|
||||
`--pooler-config.task token_embed` online.
|
||||
|
||||
## Typical Use Cases
|
||||
|
||||
@@ -400,6 +400,7 @@ th {
|
||||
| `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |
|
||||
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
|
||||
| `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
|
||||
| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
|
||||
| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -554,6 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
|
||||
| `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
|
||||
| `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
|
||||
| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
|
||||
| `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -633,6 +635,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
|
||||
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
|
||||
<sup>+</sup> Multiple items can be inputted per text prompt for this modality.
|
||||
<sup>*</sup> Only specific variants of the model support this modality (see notes below).</br>
|
||||
|
||||
!!! note
|
||||
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
|
||||
@@ -643,6 +646,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
|
||||
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
|
||||
|
||||
!!! note
|
||||
For `Gemma4ForConditionalGeneration`:
|
||||
- audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
|
||||
- The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
|
||||
|
||||
!!! note
|
||||
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
|
||||
|
||||
@@ -674,6 +682,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
|
||||
!!! note
|
||||
`VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
|
||||
|
||||
#### Realtime Transcription
|
||||
|
||||
Speech models that support streaming transcription via the
|
||||
[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
|
||||
WebSocket endpoint.
|
||||
|
||||
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
|
||||
| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
|
||||
| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
|
||||
|
||||
!!! note
|
||||
`VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
|
||||
|
||||
`Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
|
||||
You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
|
||||
when serving.
|
||||
|
||||
## Pooling Models
|
||||
|
||||
See [this page](pooling_models/README.md) for more information on how to use pooling models.
|
||||
|
||||
@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
|
||||
- [Translation API](#translations-api) (`/v1/audio/translations`)
|
||||
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
|
||||
- [Realtime API](#realtime-api) (`/v1/realtime`)
|
||||
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
|
||||
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
|
||||
|
||||
In addition, we have the following custom APIs:
|
||||
|
||||
@@ -467,28 +467,11 @@ It consists of two endpoints:
|
||||
- `/tokenize` corresponds to calling `tokenizer.encode()`.
|
||||
- `/detokenize` corresponds to calling `tokenizer.decode()`.
|
||||
|
||||
### Score API
|
||||
|
||||
#### Score Template
|
||||
|
||||
Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)).
|
||||
|
||||
Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template.
|
||||
|
||||
Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
|
||||
|
||||
- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
|
||||
- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}`
|
||||
|
||||
This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future.
|
||||
|
||||
Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
|
||||
|
||||
### Generative Scoring API
|
||||
|
||||
The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes".
|
||||
|
||||
This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
|
||||
This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
|
||||
|
||||
**Requirements:**
|
||||
|
||||
|
||||
+2
-2
@@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml
|
||||
|
||||
For detailed deployment instructions and platform-specific options, see:
|
||||
|
||||
- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import
|
||||
- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping
|
||||
- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import
|
||||
- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping
|
||||
|
||||
## Contributing
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Disaggregated multimodal serving: render → generate round-trip.
|
||||
|
||||
Demonstrates the two-phase disaggregated flow:
|
||||
1. /v1/chat/completions/render – preprocesses a multimodal chat request
|
||||
into token IDs and serialized tensor features.
|
||||
2. /inference/v1/generate – runs inference on the preprocessed tokens.
|
||||
|
||||
The render response is passed *directly* to generate with only
|
||||
``sampling_params`` added, showing that the two endpoints compose with
|
||||
zero client-side transformation.
|
||||
|
||||
Launch the server first:
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-2B-Instruct \
|
||||
--dtype bfloat16 --max-model-len 4096 --enforce-eager
|
||||
|
||||
Then run this script:
|
||||
|
||||
python example_mm_serve.py
|
||||
"""
|
||||
|
||||
import io
|
||||
|
||||
import pybase64 as base64
|
||||
import requests
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
|
||||
|
||||
|
||||
def make_data_url(image: Image.Image) -> str:
|
||||
"""Encode a PIL image as a base64 data URL."""
|
||||
buf = io.BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
|
||||
def main():
|
||||
# -- Step 1: Create a test image (solid red) -------------------------
|
||||
image = Image.new("RGB", (224, 224), color=(255, 0, 0))
|
||||
data_url = make_data_url(image)
|
||||
print("Created 224x224 red test image")
|
||||
|
||||
# -- Step 2: Render (preprocess) -------------------------------------
|
||||
render_payload = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What color is this image? Answer in one word.",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
print("\n--- Render ---")
|
||||
render_resp = requests.post(
|
||||
f"{BASE_URL}/v1/chat/completions/render", json=render_payload
|
||||
)
|
||||
render_resp.raise_for_status()
|
||||
render_data = render_resp.json()
|
||||
|
||||
print(f"Response keys: {list(render_data.keys())}")
|
||||
print(f"Number of token_ids: {len(render_data['token_ids'])}")
|
||||
|
||||
features = render_data.get("features")
|
||||
if features and features.get("kwargs_data"):
|
||||
print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
|
||||
for modality, items in features["kwargs_data"].items():
|
||||
print(
|
||||
f" {modality}: {len(items)} item(s), "
|
||||
f"first item type: {type(items[0])} length: {len(items[0])}"
|
||||
if items
|
||||
else "First item: (empty)"
|
||||
)
|
||||
else:
|
||||
print("WARNING: no kwargs_data in render response")
|
||||
|
||||
# -- Step 3: Generate (inference) ------------------------------------
|
||||
# Pass the render output directly — only add sampling_params.
|
||||
generate_payload = render_data
|
||||
generate_payload["sampling_params"] = {
|
||||
"max_tokens": 20,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
print("\n--- Generate ---")
|
||||
gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
|
||||
gen_resp.raise_for_status()
|
||||
gen_data = gen_resp.json()
|
||||
|
||||
# -- Step 4: Decode & print ------------------------------------------
|
||||
output_ids = gen_data["choices"][0]["token_ids"]
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
text = tokenizer.decode(output_ids, skip_special_tokens=True)
|
||||
|
||||
print(f"Output token count: {len(output_ids)}")
|
||||
print(f"Generated text: {text!r}")
|
||||
|
||||
if "red" in text.lower():
|
||||
print("\nModel correctly identified the red image.")
|
||||
else:
|
||||
print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -267,7 +267,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {
|
||||
# Any format supported by librosa is supported
|
||||
# Any format supported by soundfile/PyAV is supported
|
||||
"data": audio_base64,
|
||||
"format": "wav",
|
||||
},
|
||||
@@ -292,7 +292,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
# Any format supported by librosa is supported
|
||||
# Any format supported by soundfile/PyAV is supported
|
||||
"url": audio_url
|
||||
},
|
||||
},
|
||||
@@ -316,7 +316,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
# Any format supported by librosa is supported
|
||||
# Any format supported by soundfile/PyAV is supported
|
||||
"url": f"data:audio/ogg;base64,{audio_base64}"
|
||||
},
|
||||
},
|
||||
|
||||
@@ -12,7 +12,6 @@ model, for example:
|
||||
Requirements:
|
||||
- vllm with audio support
|
||||
- websockets
|
||||
- librosa
|
||||
- numpy
|
||||
|
||||
The script:
|
||||
@@ -26,12 +25,12 @@ import argparse
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import pybase64 as base64
|
||||
import websockets
|
||||
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
|
||||
def audio_to_pcm16_base64(audio_path: str) -> str:
|
||||
@@ -39,7 +38,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str:
|
||||
Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
|
||||
"""
|
||||
# Load audio and resample to 16kHz mono
|
||||
audio, _ = librosa.load(audio_path, sr=16000, mono=True)
|
||||
audio, _ = load_audio(audio_path, sr=16000, mono=True)
|
||||
# Convert to PCM16
|
||||
pcm16 = (audio * 32767).astype(np.int16)
|
||||
# Encode as base64
|
||||
|
||||
@@ -170,6 +170,7 @@ eles = "eles"
|
||||
datas = "datas"
|
||||
ser = "ser"
|
||||
ure = "ure"
|
||||
VALU = "VALU"
|
||||
# Walsh-Hadamard Transform
|
||||
wht = "wht"
|
||||
WHT = "WHT"
|
||||
|
||||
@@ -7,7 +7,7 @@ requests >= 2.26.0
|
||||
tqdm
|
||||
blake3
|
||||
py-cpuinfo
|
||||
transformers >= 4.56.0, < 5
|
||||
transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
|
||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
|
||||
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
|
||||
@@ -19,7 +19,7 @@ pillow # Required for image processing
|
||||
prometheus-fastapi-instrumentator >= 7.0.0
|
||||
tiktoken >= 0.6.0 # Required for DBRX tokenizer
|
||||
lm-format-enforcer == 0.11.3
|
||||
llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x" or platform_machine == "ppc64le"
|
||||
llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
|
||||
outlines_core == 0.2.11
|
||||
# required for outlines backend disk cache
|
||||
diskcache == 5.6.3
|
||||
@@ -37,7 +37,7 @@ pyyaml
|
||||
six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
|
||||
setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
|
||||
einops # Required for Qwen2-VL.
|
||||
compressed-tensors == 0.14.0.1 # required for compressed-tensors
|
||||
compressed-tensors == 0.15.0.1 # required for compressed-tensors
|
||||
depyf==0.20.0 # required for profiling and debugging with compilation config
|
||||
cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
|
||||
watchfiles # required for http server to monitor the updates of TLS files
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
lmcache >= 0.3.9
|
||||
nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
|
||||
nixl-cu12 >= 0.7.1, < 0.10.0
|
||||
nixl-cu13 >= 0.7.1, < 0.10.0
|
||||
nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
|
||||
nixl-cu12 >= 0.7.1, <= 0.10.1
|
||||
nixl-cu13 >= 0.7.1, <= 0.10.1
|
||||
mooncake-transfer-engine >= 0.3.8
|
||||
|
||||
@@ -20,6 +20,4 @@ conch-triton-kernels==1.2.1
|
||||
timm>=1.0.17
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
# To be consistent with test_quark.py
|
||||
amd-quark>=0.8.99
|
||||
# Required for faster safetensors model loading
|
||||
fastsafetensors >= 0.2.2
|
||||
amd-quark>=0.8.99
|
||||
@@ -18,10 +18,9 @@ httpx
|
||||
librosa # required for audio tests
|
||||
vector_quantize_pytorch # required for minicpmo_26 test
|
||||
vocos # required for minicpmo_26 test
|
||||
peft>=0.15.0 # required for phi-4-mm test
|
||||
peft>=0.18.1 # required for phi-4-mm test
|
||||
pqdm
|
||||
ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
|
||||
resampy # required for audio tests
|
||||
sentence-transformers>=5.2.0 # required for embedding tests
|
||||
soundfile # required for audio tests
|
||||
jiwer # required for audio tests
|
||||
@@ -39,8 +38,8 @@ opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
transformers==5.5.3
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=3.39.15 # Required for openai schema test.
|
||||
# quantization
|
||||
bitsandbytes==0.49.2
|
||||
|
||||
+10
-14
@@ -4,7 +4,7 @@ absl-py==2.1.0
|
||||
# via
|
||||
# rouge-score
|
||||
# tensorboard
|
||||
accelerate==1.0.1
|
||||
accelerate==1.13.0
|
||||
# via peft
|
||||
aenum==3.1.16
|
||||
# via lightly
|
||||
@@ -248,7 +248,6 @@ filelock==3.16.1
|
||||
# huggingface-hub
|
||||
# ray
|
||||
# torch
|
||||
# transformers
|
||||
# virtualenv
|
||||
fiona==1.10.1
|
||||
# via torchgeo
|
||||
@@ -331,7 +330,7 @@ h5py==3.13.0
|
||||
# via terratorch
|
||||
harfile==0.3.0
|
||||
# via schemathesis
|
||||
hf-xet==1.1.7
|
||||
hf-xet==1.4.3
|
||||
# via huggingface-hub
|
||||
hiredis==3.0.0
|
||||
# via tensorizer
|
||||
@@ -345,9 +344,10 @@ httpx==0.27.2
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
# diffusers
|
||||
# huggingface-hub
|
||||
# perceptron
|
||||
# schemathesis
|
||||
huggingface-hub==0.36.2
|
||||
huggingface-hub==1.10.2
|
||||
# via
|
||||
# accelerate
|
||||
# datasets
|
||||
@@ -555,7 +555,6 @@ numba==0.61.2
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
# librosa
|
||||
# resampy
|
||||
numpy==2.2.6
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -596,7 +595,6 @@ numpy==2.2.6
|
||||
# pyogrio
|
||||
# pywavelets
|
||||
# rasterio
|
||||
# resampy
|
||||
# rioxarray
|
||||
# rouge-score
|
||||
# runai-model-streamer
|
||||
@@ -756,7 +754,7 @@ pathvalidate==3.2.1
|
||||
# via pytablewriter
|
||||
patsy==1.0.1
|
||||
# via statsmodels
|
||||
peft==0.16.0
|
||||
peft==0.18.1
|
||||
# via -r requirements/test/cuda.in
|
||||
perceptron==0.1.4
|
||||
# via -r requirements/test/cuda.in
|
||||
@@ -982,7 +980,7 @@ referencing==0.35.1
|
||||
# via
|
||||
# jsonschema
|
||||
# jsonschema-specifications
|
||||
regex==2024.9.11
|
||||
regex==2026.2.28
|
||||
# via
|
||||
# diffusers
|
||||
# nltk
|
||||
@@ -1002,7 +1000,6 @@ requests==2.32.3
|
||||
# google-api-core
|
||||
# google-cloud-storage
|
||||
# gpt-oss
|
||||
# huggingface-hub
|
||||
# lightly
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
@@ -1015,10 +1012,7 @@ requests==2.32.3
|
||||
# starlette-testclient
|
||||
# tacoreader
|
||||
# tiktoken
|
||||
# transformers
|
||||
# wandb
|
||||
resampy==0.4.3
|
||||
# via -r requirements/test/cuda.in
|
||||
responses==0.25.3
|
||||
# via genai-perf
|
||||
rfc3339-validator==0.1.4
|
||||
@@ -1216,7 +1210,7 @@ timm==1.0.17
|
||||
# segmentation-models-pytorch
|
||||
# terratorch
|
||||
# torchgeo
|
||||
tokenizers==0.22.0
|
||||
tokenizers==0.22.2
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1295,7 +1289,7 @@ tqdm==4.67.3
|
||||
# tacoreader
|
||||
# terratorch
|
||||
# transformers
|
||||
transformers==4.57.5
|
||||
transformers==5.5.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1317,7 +1311,9 @@ typepy==1.3.2
|
||||
typer==0.15.2
|
||||
# via
|
||||
# fastsafetensors
|
||||
# huggingface-hub
|
||||
# perceptron
|
||||
# transformers
|
||||
types-python-dateutil==2.9.0.20241206
|
||||
# via arrow
|
||||
typeshed-client==2.8.2
|
||||
|
||||
@@ -29,8 +29,8 @@ opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
transformers==5.5.3
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=3.39.15 # Required for openai schema test.
|
||||
# quantization
|
||||
bitsandbytes>=0.49.2
|
||||
|
||||
@@ -23,7 +23,6 @@ vocos # required for minicpmo_26 test
|
||||
peft>=0.15.0 # required for phi-4-mm test
|
||||
pqdm
|
||||
ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
|
||||
resampy # required for audio tests
|
||||
sentence-transformers>=5.2.0 # required for embedding tests
|
||||
soundfile # required for audio tests
|
||||
jiwer # required for audio tests
|
||||
@@ -38,8 +37,8 @@ opencv-python-headless>=4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
transformers==5.5.3
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=3.39.15 # Required for openai schema test
|
||||
# quantization
|
||||
bitsandbytes==0.49.2
|
||||
@@ -56,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
|
||||
numba==0.61.2 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
decord==0.6.0
|
||||
@@ -82,4 +81,3 @@ plotly # required for perf comparison html report
|
||||
rapidfuzz
|
||||
torchgeo==0.7.0
|
||||
multiprocess==0.70.16
|
||||
huggingface-hub==0.36.2
|
||||
|
||||
+17
-24
@@ -39,7 +39,7 @@ annotated-doc==0.0.4
|
||||
# typer
|
||||
annotated-types==0.7.0
|
||||
# via pydantic
|
||||
anthropic==0.89.0
|
||||
anthropic==0.93.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -172,7 +172,7 @@ colorful==0.5.8
|
||||
# via ray
|
||||
colorlog==6.10.1
|
||||
# via optuna
|
||||
compressed-tensors==0.14.0.1
|
||||
compressed-tensors==0.15.0.1
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -269,16 +269,14 @@ fastapi==0.135.2
|
||||
# model-hosting-container-standards
|
||||
fastapi-cli==0.0.24
|
||||
# via fastapi
|
||||
fastapi-cloud-cli==0.15.1
|
||||
fastapi-cloud-cli==0.16.1
|
||||
# via fastapi-cli
|
||||
fastar==0.9.0
|
||||
fastar==0.10.0
|
||||
# via fastapi-cloud-cli
|
||||
fastparquet==2026.3.0
|
||||
# via genai-perf
|
||||
fastsafetensors==0.2.2
|
||||
# via
|
||||
# -c requirements/rocm.txt
|
||||
# -r requirements/test/rocm.in
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
|
||||
# via -r requirements/test/rocm.in
|
||||
filelock==3.25.2
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
@@ -290,7 +288,6 @@ filelock==3.25.2
|
||||
# python-discovery
|
||||
# ray
|
||||
# torch
|
||||
# transformers
|
||||
# virtualenv
|
||||
fiona==1.10.1
|
||||
# via torchgeo
|
||||
@@ -384,7 +381,7 @@ h5py==3.16.0
|
||||
# via terratorch
|
||||
harfile==0.4.0
|
||||
# via schemathesis
|
||||
hf-xet==1.4.2
|
||||
hf-xet==1.4.3
|
||||
# via huggingface-hub
|
||||
hiredis==3.3.1
|
||||
# via tensorizer
|
||||
@@ -403,6 +400,7 @@ httpx==0.27.2
|
||||
# diffusers
|
||||
# fastapi
|
||||
# fastapi-cloud-cli
|
||||
# huggingface-hub
|
||||
# mcp
|
||||
# model-hosting-container-standards
|
||||
# openai
|
||||
@@ -410,9 +408,8 @@ httpx==0.27.2
|
||||
# schemathesis
|
||||
httpx-sse==0.4.3
|
||||
# via mcp
|
||||
huggingface-hub==0.36.2
|
||||
huggingface-hub==1.10.2
|
||||
# via
|
||||
# -r requirements/test/rocm.in
|
||||
# accelerate
|
||||
# datasets
|
||||
# diffusers
|
||||
@@ -484,7 +481,7 @@ jinja2==3.1.6
|
||||
# genai-perf
|
||||
# lm-eval
|
||||
# torch
|
||||
jiter==0.13.0
|
||||
jiter==0.14.0
|
||||
# via
|
||||
# anthropic
|
||||
# openai
|
||||
@@ -631,7 +628,7 @@ msgpack==1.1.2
|
||||
# via
|
||||
# librosa
|
||||
# ray
|
||||
msgspec==0.20.0
|
||||
msgspec==0.21.0
|
||||
# via -r requirements/test/../common.txt
|
||||
mteb==2.11.5
|
||||
# via -r requirements/test/rocm.in
|
||||
@@ -663,7 +660,6 @@ numba==0.61.2
|
||||
# -c requirements/rocm.txt
|
||||
# -r requirements/test/rocm.in
|
||||
# librosa
|
||||
# resampy
|
||||
numkong==7.1.1
|
||||
# via albucore
|
||||
numpy==2.2.6
|
||||
@@ -709,7 +705,6 @@ numpy==2.2.6
|
||||
# pytrec-eval-terrier
|
||||
# pywavelets
|
||||
# rasterio
|
||||
# resampy
|
||||
# rioxarray
|
||||
# rouge-score
|
||||
# runai-model-streamer
|
||||
@@ -742,7 +737,7 @@ omegaconf==2.3.0
|
||||
# lightning
|
||||
open-clip-torch==2.32.0
|
||||
# via -r requirements/test/rocm.in
|
||||
openai==2.30.0
|
||||
openai==2.31.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -1093,7 +1088,7 @@ python-dotenv==1.2.2
|
||||
# uvicorn
|
||||
python-json-logger==4.1.0
|
||||
# via -r requirements/test/../common.txt
|
||||
python-multipart==0.0.22
|
||||
python-multipart==0.0.26
|
||||
# via
|
||||
# fastapi
|
||||
# mcp
|
||||
@@ -1180,7 +1175,6 @@ requests==2.32.5
|
||||
# google-api-core
|
||||
# google-cloud-storage
|
||||
# gpt-oss
|
||||
# huggingface-hub
|
||||
# lightly
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
@@ -1194,10 +1188,7 @@ requests==2.32.5
|
||||
# starlette-testclient
|
||||
# tacoreader
|
||||
# tiktoken
|
||||
# transformers
|
||||
# wandb
|
||||
resampy==0.4.3
|
||||
# via -r requirements/test/rocm.in
|
||||
responses==0.26.0
|
||||
# via genai-perf
|
||||
rfc3339-validator==0.1.4
|
||||
@@ -1428,7 +1419,7 @@ timm==1.0.17
|
||||
# segmentation-models-pytorch
|
||||
# terratorch
|
||||
# torchgeo
|
||||
tokenizers==0.22.0
|
||||
tokenizers==0.22.2
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -1471,7 +1462,7 @@ tqdm==4.67.3
|
||||
# tacoreader
|
||||
# terratorch
|
||||
# transformers
|
||||
transformers==4.57.5
|
||||
transformers==5.5.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -1498,7 +1489,9 @@ typer==0.24.1
|
||||
# fastapi-cli
|
||||
# fastapi-cloud-cli
|
||||
# fastsafetensors
|
||||
# huggingface-hub
|
||||
# perceptron
|
||||
# transformers
|
||||
typeshed-client==2.9.0
|
||||
# via jsonargparse
|
||||
typing-extensions==4.15.0
|
||||
|
||||
@@ -13,7 +13,6 @@ pytest-shard
|
||||
absl-py
|
||||
accelerate
|
||||
arctic-inference
|
||||
hf_transfer
|
||||
lm_eval[api]
|
||||
modelscope
|
||||
|
||||
|
||||
@@ -19,7 +19,9 @@ aiosignal==1.4.0
|
||||
albumentations==1.4.6
|
||||
# via -r requirements/test/xpu.in
|
||||
annotated-doc==0.0.4
|
||||
# via fastapi
|
||||
# via
|
||||
# fastapi
|
||||
# typer
|
||||
annotated-types==0.7.0
|
||||
# via pydantic
|
||||
anyio==4.13.0
|
||||
@@ -64,6 +66,7 @@ click==8.3.1
|
||||
# jiwer
|
||||
# nltk
|
||||
# schemathesis
|
||||
# typer
|
||||
# uvicorn
|
||||
colorama==0.4.6
|
||||
# via sacrebleu
|
||||
@@ -112,7 +115,6 @@ filelock==3.25.2
|
||||
# huggingface-hub
|
||||
# modelscope
|
||||
# torch
|
||||
# transformers
|
||||
frozenlist==1.8.0
|
||||
# via
|
||||
# aiohttp
|
||||
@@ -133,9 +135,7 @@ h11==0.16.0
|
||||
# uvicorn
|
||||
harfile==0.4.0
|
||||
# via schemathesis
|
||||
hf-transfer==0.1.9
|
||||
# via -r requirements/test/xpu.in
|
||||
hf-xet==1.4.2
|
||||
hf-xet==1.4.3
|
||||
# via huggingface-hub
|
||||
html2text==2025.4.15
|
||||
# via gpt-oss
|
||||
@@ -144,8 +144,9 @@ httpcore==1.0.9
|
||||
httpx==0.28.1
|
||||
# via
|
||||
# datasets
|
||||
# huggingface-hub
|
||||
# schemathesis
|
||||
huggingface-hub==0.36.2
|
||||
huggingface-hub==1.10.2
|
||||
# via
|
||||
# accelerate
|
||||
# datasets
|
||||
@@ -515,7 +516,6 @@ requests==2.33.1
|
||||
# docker
|
||||
# evaluate
|
||||
# gpt-oss
|
||||
# huggingface-hub
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
# modelscope
|
||||
@@ -524,11 +524,11 @@ requests==2.33.1
|
||||
# schemathesis
|
||||
# starlette-testclient
|
||||
# tiktoken
|
||||
# transformers
|
||||
rich==14.3.3
|
||||
# via
|
||||
# mteb
|
||||
# schemathesis
|
||||
# typer
|
||||
rouge-score==0.1.2
|
||||
# via lm-eval
|
||||
rpds-py==0.30.0
|
||||
@@ -572,6 +572,8 @@ setuptools==80.10.2
|
||||
# modelscope
|
||||
# pytablewriter
|
||||
# torch
|
||||
shellingham==1.5.4
|
||||
# via typer
|
||||
six==1.17.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
@@ -665,7 +667,7 @@ tqdm==4.67.3
|
||||
# pqdm
|
||||
# sentence-transformers
|
||||
# transformers
|
||||
transformers==4.57.6
|
||||
transformers==5.5.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# sentence-transformers
|
||||
@@ -676,6 +678,10 @@ typepy==1.3.4
|
||||
# dataproperty
|
||||
# pytablewriter
|
||||
# tabledata
|
||||
typer==0.24.1
|
||||
# via
|
||||
# huggingface-hub
|
||||
# transformers
|
||||
typing-extensions==4.15.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
|
||||
@@ -693,6 +693,12 @@ class precompiled_wheel_utils:
|
||||
flash_attn_regex = re.compile(
|
||||
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
|
||||
)
|
||||
# __init__.py and flash_attn_interface.py are source-controlled
|
||||
# in vllm and should not be overwritten (matches cmake exclusions)
|
||||
flash_attn_files_to_skip = {
|
||||
"vllm/vllm_flash_attn/__init__.py",
|
||||
"vllm/vllm_flash_attn/flash_attn_interface.py",
|
||||
}
|
||||
triton_kernels_regex = re.compile(
|
||||
r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
|
||||
)
|
||||
@@ -705,7 +711,11 @@ class precompiled_wheel_utils:
|
||||
filter(lambda x: x.filename in files_to_copy, wheel.filelist)
|
||||
)
|
||||
file_members += list(
|
||||
filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
|
||||
filter(
|
||||
lambda x: flash_attn_regex.match(x.filename)
|
||||
and x.filename not in flash_attn_files_to_skip,
|
||||
wheel.filelist,
|
||||
)
|
||||
)
|
||||
file_members += list(
|
||||
filter(
|
||||
@@ -1075,7 +1085,9 @@ setup(
|
||||
install_requires=get_requirements(),
|
||||
extras_require={
|
||||
# AMD Zen CPU optimizations via zentorch
|
||||
"zen": ["zentorch"],
|
||||
"zen": [
|
||||
"zentorch-weekly==5.2.1.dev20260408"
|
||||
], # Zentorch has weekly releases. This pulls the known-good version.
|
||||
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
|
||||
"tensorizer": ["tensorizer==2.10.1"],
|
||||
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
|
||||
@@ -1083,7 +1095,6 @@ setup(
|
||||
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
|
||||
"audio": [
|
||||
"av",
|
||||
"resampy",
|
||||
"scipy",
|
||||
"soundfile",
|
||||
"mistral_common[audio]",
|
||||
@@ -1094,7 +1105,7 @@ setup(
|
||||
# NOTE: When updating helion version, also update CI files:
|
||||
# - .buildkite/test_areas/kernels.yaml
|
||||
# - .buildkite/test-amd.yaml
|
||||
"helion": ["helion==0.3.3"],
|
||||
"helion": ["helion==1.0.0"],
|
||||
# Optional deps for gRPC server (vllm serve --grpc)
|
||||
"grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"],
|
||||
# Optional deps for OpenTelemetry tracing
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.config
|
||||
from tests.compile.backend import TestBackend
|
||||
from vllm.compilation.passes.vllm_inductor_pass import (
|
||||
VllmFusionPatternMatcherPass,
|
||||
VllmPatternMatcherPass,
|
||||
VllmPatternReplacement,
|
||||
)
|
||||
from vllm.config import CompilationConfig, CompilationMode, VllmConfig
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
class ReluToAbsPattern(VllmPatternReplacement):
|
||||
"""Replaces relu(x) with abs(x) — a minimal test fixture."""
|
||||
|
||||
@property
|
||||
def pattern(self):
|
||||
def _pattern(x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.aten.relu.default(x)
|
||||
|
||||
return _pattern
|
||||
|
||||
@property
|
||||
def replacement(self):
|
||||
def _replacement(x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.aten.abs.default(x)
|
||||
|
||||
return _replacement
|
||||
|
||||
def get_inputs(self) -> list[torch.Tensor]:
|
||||
return [self.empty_fp32(4)]
|
||||
|
||||
|
||||
class ExpToSqrtPattern(VllmPatternReplacement):
|
||||
"""A second distinct pattern type — used to test uuid differentiation."""
|
||||
|
||||
@property
|
||||
def pattern(self):
|
||||
def _pattern(x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.aten.exp.default(x)
|
||||
|
||||
return _pattern
|
||||
|
||||
@property
|
||||
def replacement(self):
|
||||
def _replacement(x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.aten.sqrt.default(x)
|
||||
|
||||
return _replacement
|
||||
|
||||
def get_inputs(self) -> list[torch.Tensor]:
|
||||
return [self.empty_fp32(4)]
|
||||
|
||||
|
||||
class ReluFusionPass(VllmFusionPatternMatcherPass):
|
||||
def __init__(self, config: VllmConfig) -> None:
|
||||
super().__init__(config, "test_relu_fusion")
|
||||
self.register(ReluToAbsPattern())
|
||||
|
||||
|
||||
class TwoPatternFusionPass(VllmFusionPatternMatcherPass):
|
||||
def __init__(self, config: VllmConfig) -> None:
|
||||
super().__init__(config, "test_two_pattern_fusion")
|
||||
self.register(ReluToAbsPattern())
|
||||
self.register(ExpToSqrtPattern())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vllm_config():
|
||||
return VllmConfig(
|
||||
compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
|
||||
def test_register_tracks_patterns(vllm_config):
|
||||
"""register() appends each VllmPatternReplacement to _pattern_replacements."""
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
single = ReluFusionPass(vllm_config)
|
||||
two = TwoPatternFusionPass(vllm_config)
|
||||
|
||||
assert len(single._pattern_replacements) == 1
|
||||
assert len(two._pattern_replacements) == 2
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
|
||||
def test_uuid_stable(vllm_config):
|
||||
"""Two instances of the same pass class produce identical uuids."""
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
p1 = ReluFusionPass(vllm_config)
|
||||
p2 = ReluFusionPass(vllm_config)
|
||||
p3 = TwoPatternFusionPass(vllm_config)
|
||||
|
||||
assert p1.uuid() == p2.uuid()
|
||||
assert p1.uuid() != p3.uuid()
|
||||
assert p2.uuid() != p3.uuid()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
|
||||
@pytest.mark.parametrize("N", [1, 2, 4])
|
||||
def test_matched_count_and_match_table(vllm_config, N):
|
||||
"""matched_count and match_table reflect the number of matched patterns."""
|
||||
|
||||
class Model(torch.nn.Module):
|
||||
def forward(self, *inputs):
|
||||
# N independent relus
|
||||
return sum(torch.relu(x) for x in inputs)
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(torch.float32)
|
||||
|
||||
fusion_pass = ReluFusionPass(vllm_config)
|
||||
backend = TestBackend(fusion_pass)
|
||||
model = torch.compile(Model(), backend=backend)
|
||||
|
||||
inputs = [torch.rand(8) for _ in range(N)]
|
||||
model(*inputs)
|
||||
|
||||
assert fusion_pass.matched_count == N
|
||||
assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N
|
||||
@@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards(
|
||||
torch.randn(1, 10).cuda(),
|
||||
is_01_specialization=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
|
||||
def test_piecewise_backend_empty_sym_shape_indices():
|
||||
"""Test that PiecewiseBackend handles empty sym_shape_indices correctly.
|
||||
|
||||
When all inputs have static shapes (no torch.SymInt), sym_shape_indices
|
||||
will be empty. The fix in PiecewiseBackend.__call__ handles this case
|
||||
by using the first compiled range_entry.
|
||||
"""
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Use small max_model_len and max_num_batched_tokens to encourage
|
||||
# static shape compilation with empty sym_shape_indices
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
max_model_len=512,
|
||||
max_num_batched_tokens=1,
|
||||
compilation_config={
|
||||
"mode": CompilationMode.VLLM_COMPILE,
|
||||
"dynamic_shapes_config": {
|
||||
"type": DynamicShapesType.BACKED.value,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
|
||||
|
||||
# Generate with static shape inputs
|
||||
output = llm.generate("Hello, my name is", sampling_params=sampling_params)
|
||||
result = output[0].outputs[0].text
|
||||
assert len(result) > 0, "Should generate non-empty output"
|
||||
|
||||
# Generate again to verify compilation works with empty sym_shape_indices
|
||||
output = llm.generate("The capital of France is", sampling_params=sampling_params)
|
||||
result = output[0].outputs[0].text
|
||||
assert len(result) > 0, "Should generate non-empty output on second run"
|
||||
|
||||
del llm
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
@@ -356,6 +356,23 @@
|
||||
"is_multimodal_model": false,
|
||||
"dtype": "torch.float32"
|
||||
},
|
||||
"stepfun-ai/Step-3.5-Flash": {
|
||||
"architectures": [
|
||||
"Step3p5ForCausalLM"
|
||||
],
|
||||
"model_type": "step3p5",
|
||||
"text_model_type": "step3p5",
|
||||
"hidden_size": 4096,
|
||||
"total_num_hidden_layers": 45,
|
||||
"total_num_attention_heads": 64,
|
||||
"head_size": 128,
|
||||
"vocab_size": 128896,
|
||||
"total_num_kv_heads": 8,
|
||||
"num_experts": 288,
|
||||
"is_deepseek_mla": false,
|
||||
"is_multimodal_model": false,
|
||||
"dtype": "torch.bfloat16"
|
||||
},
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": {
|
||||
"architectures": [
|
||||
"NemotronHForCausalLM"
|
||||
|
||||
@@ -16,6 +16,7 @@ BASE_TRUST_REMOTE_CODE_MODELS = {
|
||||
"nvidia/Llama-3_3-Nemotron-Super-49B-v1",
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
|
||||
"XiaomiMiMo/MiMo-7B-RL",
|
||||
"stepfun-ai/Step-3.5-Flash",
|
||||
# Excluded: Not available online right now
|
||||
# "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",
|
||||
"meituan-longcat/LongCat-Flash-Chat",
|
||||
|
||||
@@ -364,6 +364,7 @@ class HfRunner:
|
||||
model_name: str,
|
||||
dtype: str = "auto",
|
||||
*,
|
||||
revision: str | None = None,
|
||||
model_kwargs: dict[str, Any] | None = None,
|
||||
trust_remote_code: bool = True,
|
||||
is_sentence_transformer: bool = False,
|
||||
@@ -383,6 +384,7 @@ class HfRunner:
|
||||
self._init(
|
||||
model_name=model_name,
|
||||
dtype=dtype,
|
||||
revision=revision,
|
||||
model_kwargs=model_kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
is_sentence_transformer=is_sentence_transformer,
|
||||
@@ -396,6 +398,7 @@ class HfRunner:
|
||||
model_name: str,
|
||||
dtype: str = "auto",
|
||||
*,
|
||||
revision: str | None = None,
|
||||
model_kwargs: dict[str, Any] | None = None,
|
||||
trust_remote_code: bool = True,
|
||||
is_sentence_transformer: bool = False,
|
||||
@@ -410,6 +413,15 @@ class HfRunner:
|
||||
model_name,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
# HF runner should use the HF config so that it's consistent with the HF model
|
||||
if self.config.__module__.startswith("vllm.transformers_utils.configs"):
|
||||
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
|
||||
|
||||
del CONFIG_MAPPING._extra_content[self.config.model_type]
|
||||
self.config = AutoConfig.from_pretrained(
|
||||
model_name,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
self.device = self.get_default_device()
|
||||
self.dtype = dtype = _get_and_verify_dtype(
|
||||
self.model_name,
|
||||
@@ -428,6 +440,7 @@ class HfRunner:
|
||||
|
||||
self.model = SentenceTransformer(
|
||||
model_name,
|
||||
revision=revision,
|
||||
device=self.device,
|
||||
model_kwargs=model_kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
@@ -438,6 +451,7 @@ class HfRunner:
|
||||
|
||||
self.model = CrossEncoder(
|
||||
model_name,
|
||||
revision=revision,
|
||||
device=self.device,
|
||||
automodel_args=model_kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
@@ -447,6 +461,7 @@ class HfRunner:
|
||||
nn.Module,
|
||||
auto_cls.from_pretrained(
|
||||
model_name,
|
||||
revision=revision,
|
||||
trust_remote_code=trust_remote_code,
|
||||
**model_kwargs,
|
||||
),
|
||||
|
||||
@@ -38,6 +38,8 @@ llm = LLM(
|
||||
distributed_executor_backend="external_launcher",
|
||||
gpu_memory_utilization=random.uniform(0.7, 0.9),
|
||||
seed=0,
|
||||
max_model_len=1024,
|
||||
max_num_seqs=16,
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
@@ -41,6 +41,7 @@ def create_mock_parallel_config(
|
||||
config.rank = rank
|
||||
config.world_size = world_size
|
||||
config.data_parallel_rank = dp_rank
|
||||
config.data_parallel_index = dp_rank
|
||||
return config
|
||||
|
||||
|
||||
@@ -283,6 +284,7 @@ def inference_receive_tensor(
|
||||
parallel_config.rank = 0
|
||||
parallel_config.world_size = 1
|
||||
parallel_config.data_parallel_rank = 0
|
||||
parallel_config.data_parallel_index = 0
|
||||
|
||||
engine = NCCLWeightTransferEngine(config, parallel_config)
|
||||
|
||||
@@ -666,6 +668,7 @@ def inference_receive_ipc_tensor(
|
||||
parallel_config.rank = 0
|
||||
parallel_config.world_size = 1
|
||||
parallel_config.data_parallel_rank = 0
|
||||
parallel_config.data_parallel_index = 0
|
||||
|
||||
engine = IPCWeightTransferEngine(config, parallel_config)
|
||||
|
||||
|
||||
@@ -91,6 +91,12 @@ def test_multiple_priority(llm: LLM):
|
||||
outputs = llm.generate(PROMPTS, sampling_params=None, priority=[])
|
||||
|
||||
|
||||
def test_single_prompt_priority(llm: LLM):
|
||||
# Single string prompts should be normalized to one request.
|
||||
outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=[0])
|
||||
assert len(outputs) == 1
|
||||
|
||||
|
||||
def test_max_model_len():
|
||||
max_model_len = 20
|
||||
llm = LLM(
|
||||
|
||||
@@ -13,7 +13,6 @@ import io
|
||||
import time
|
||||
from statistics import mean, median
|
||||
|
||||
import librosa
|
||||
import pytest
|
||||
import soundfile
|
||||
import torch
|
||||
@@ -21,6 +20,7 @@ from datasets import load_dataset
|
||||
from evaluate import load
|
||||
from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
|
||||
|
||||
from vllm.multimodal.audio import get_audio_duration
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
from ....models.registry import HF_EXAMPLE_MODELS
|
||||
@@ -84,7 +84,7 @@ async def process_dataset(model, client, data, concurrent_request):
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
)
|
||||
|
||||
# Warmup call as the first `librosa.load` server-side is quite slow.
|
||||
# Warmup call as the first `load_audio` server-side is quite slow.
|
||||
audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
|
||||
_ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "")
|
||||
|
||||
@@ -118,7 +118,7 @@ def print_performance_metrics(results, total_time):
|
||||
|
||||
def add_duration(sample):
|
||||
y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
|
||||
sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000
|
||||
sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000
|
||||
return sample
|
||||
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@ import asyncio
|
||||
import json
|
||||
import warnings
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import pybase64 as base64
|
||||
import pytest
|
||||
@@ -14,6 +13,7 @@ import websockets
|
||||
from tests.entrypoints.openai.conftest import add_attention_backend
|
||||
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
# Increase engine iteration timeout for ROCm where first-use JIT compilation
|
||||
# can exceed the default 60s, causing a silent deadlock in feed_tokens.
|
||||
@@ -56,7 +56,7 @@ async def send_event(ws, event: dict) -> None:
|
||||
def mary_had_lamb_audio_chunks() -> list[str]:
|
||||
"""Audio split into ~1 second chunks for streaming."""
|
||||
path = AudioAsset("mary_had_lamb").get_local_path()
|
||||
audio, _ = librosa.load(str(path), sr=16000, mono=True)
|
||||
audio, _ = load_audio(str(path), sr=16000, mono=True)
|
||||
|
||||
# Split into ~0.1 second chunks (1600 samples at 16kHz)
|
||||
chunk_size = 1600
|
||||
|
||||
@@ -6,7 +6,6 @@ import asyncio
|
||||
import io
|
||||
import json
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import openai
|
||||
import pytest
|
||||
@@ -14,6 +13,7 @@ import pytest_asyncio
|
||||
import soundfile as sf
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "openai/whisper-large-v3-turbo"
|
||||
@@ -134,7 +134,7 @@ async def test_bad_requests(mary_had_lamb, whisper_client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_long_audio_request(mary_had_lamb, whisper_client):
|
||||
mary_had_lamb.seek(0)
|
||||
audio, sr = librosa.load(mary_had_lamb)
|
||||
audio, sr = load_audio(mary_had_lamb)
|
||||
# Add small silence after each audio for repeatability in the split process
|
||||
audio = np.pad(audio, (0, 1600))
|
||||
repeated_audio = np.tile(audio, 10)
|
||||
|
||||
@@ -7,7 +7,6 @@ import io
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import librosa
|
||||
import numpy as np
|
||||
import openai
|
||||
import pytest
|
||||
@@ -17,6 +16,7 @@ import soundfile as sf
|
||||
from tests.entrypoints.openai.conftest import add_attention_backend
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.logger import init_logger
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -264,7 +264,7 @@ async def test_long_audio_request(foscolo, client_and_model):
|
||||
if model_name == "google/gemma-3n-E2B-it":
|
||||
pytest.skip("Gemma3n does not support long audio requests")
|
||||
foscolo.seek(0)
|
||||
audio, sr = librosa.load(foscolo)
|
||||
audio, sr = load_audio(foscolo)
|
||||
repeated_audio = np.tile(audio, 2)
|
||||
# Repeated audio to buffer
|
||||
buffer = io.BytesIO()
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.entrypoints.serve.disagg.mm_serde import (
|
||||
decode_mm_kwargs_item,
|
||||
encode_mm_kwargs_item,
|
||||
)
|
||||
from vllm.entrypoints.serve.disagg.protocol import (
|
||||
MultiModalFeatures,
|
||||
PlaceholderRangeInfo,
|
||||
)
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalBatchedField,
|
||||
MultiModalFieldElem,
|
||||
MultiModalFlatField,
|
||||
MultiModalKwargsItem,
|
||||
MultiModalSharedField,
|
||||
)
|
||||
|
||||
|
||||
def test_mm_kwargs_item_roundtrip():
|
||||
"""Full roundtrip test with all three field types and multiple dtypes."""
|
||||
e1 = MultiModalFieldElem(
|
||||
data=torch.zeros(1000, dtype=torch.bfloat16),
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
e2 = MultiModalFieldElem(
|
||||
data=torch.ones(100, dtype=torch.int32),
|
||||
field=MultiModalSharedField(batch_size=4),
|
||||
)
|
||||
e3 = MultiModalFieldElem(
|
||||
data=torch.randn(20, dtype=torch.float32),
|
||||
field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
|
||||
)
|
||||
|
||||
item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
|
||||
# Encoded result is a base64 string
|
||||
assert isinstance(encoded, str)
|
||||
|
||||
decoded = decode_mm_kwargs_item(encoded)
|
||||
|
||||
assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
|
||||
assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
|
||||
assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
|
||||
assert torch.equal(item["embeds"].data, decoded["embeds"].data)
|
||||
assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
|
||||
assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
|
||||
assert isinstance(decoded["embeds"].field, MultiModalFlatField)
|
||||
|
||||
|
||||
def test_mm_kwargs_item_none_data():
|
||||
"""Roundtrip with None data field."""
|
||||
elem = MultiModalFieldElem(
|
||||
data=None,
|
||||
field=MultiModalSharedField(batch_size=2),
|
||||
)
|
||||
item = MultiModalKwargsItem({"empty": elem})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
decoded = decode_mm_kwargs_item(encoded)
|
||||
|
||||
assert decoded["empty"].data is None
|
||||
assert isinstance(decoded["empty"].field, MultiModalSharedField)
|
||||
|
||||
|
||||
def test_mm_kwargs_item_nested_tensors():
|
||||
"""Roundtrip with nested tensor data."""
|
||||
nested = [torch.randn(3, 4), torch.randn(5, 4)]
|
||||
elem = MultiModalFieldElem(
|
||||
data=nested,
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
item = MultiModalKwargsItem({"nested": elem})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
decoded = decode_mm_kwargs_item(encoded)
|
||||
|
||||
decoded_data = decoded["nested"].data
|
||||
assert len(decoded_data) == 2
|
||||
assert torch.equal(nested[0], decoded_data[0])
|
||||
assert torch.equal(nested[1], decoded_data[1])
|
||||
|
||||
|
||||
def test_mm_features_with_kwargs_data():
|
||||
"""Test that MultiModalFeatures can carry serialized tensor data."""
|
||||
elem = MultiModalFieldElem(
|
||||
data=torch.randn(5, 3, dtype=torch.float32),
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
item = MultiModalKwargsItem({"pixel_values": elem})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
|
||||
features = MultiModalFeatures(
|
||||
mm_hashes={"image": ["abc123"]},
|
||||
mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
|
||||
kwargs_data={"image": [encoded]},
|
||||
)
|
||||
|
||||
# JSON roundtrip
|
||||
json_str = features.model_dump_json()
|
||||
features2 = MultiModalFeatures.model_validate_json(json_str)
|
||||
|
||||
assert features2.mm_hashes == {"image": ["abc123"]}
|
||||
assert features2.kwargs_data is not None
|
||||
assert len(features2.kwargs_data["image"]) == 1
|
||||
|
||||
decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
|
||||
assert torch.equal(elem.data, decoded["pixel_values"].data)
|
||||
@@ -0,0 +1,150 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for tool_calls Iterable → list materialisation.
|
||||
|
||||
Regression tests for https://github.com/vllm-project/vllm/issues/34792.
|
||||
|
||||
Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral
|
||||
models because:
|
||||
1. The OpenAI Python SDK types tool_calls as Iterable[...] in
|
||||
ChatCompletionAssistantMessageParam.
|
||||
2. Pydantic v2, when validating from Python objects (not from raw JSON),
|
||||
wraps Iterable fields in a one-shot lazy iterator.
|
||||
3. Debug logging called model_dump_json() which consumed that iterator.
|
||||
4. The Mistral tokenizer then saw empty tool_calls and raised
|
||||
"ValueError: Unexpected tool call id ...".
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
|
||||
|
||||
def _make_tool_call(tc_id: str, name: str, args: str) -> dict:
|
||||
return {
|
||||
"id": tc_id,
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": args},
|
||||
}
|
||||
|
||||
|
||||
def _make_request(messages: list) -> ChatCompletionRequest:
|
||||
return ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
|
||||
def test_tool_calls_list_preserved_after_model_dump():
|
||||
"""tool_calls in assistant messages must be readable after model_dump_json.
|
||||
|
||||
When the request is built from Python dicts (as in the Anthropic → OpenAI
|
||||
conversion path), Pydantic v2 previously wrapped the Iterable tool_calls
|
||||
in a one-shot iterator. model_dump_json() consumed it, leaving subsequent
|
||||
readers (e.g. the Mistral tokenizer) with an empty sequence.
|
||||
"""
|
||||
tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}')
|
||||
messages = [
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_abc123",
|
||||
"content": '{"temperature": 20}',
|
||||
},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
|
||||
# Simulate debug logging: serialize the model (this was the trigger)
|
||||
_ = req.model_dump_json()
|
||||
|
||||
# The assistant message must still have accessible tool_calls afterwards
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
tool_calls = assistant_msg.get("tool_calls")
|
||||
assert tool_calls is not None, "tool_calls must not be None after model_dump_json"
|
||||
assert isinstance(tool_calls, list), "tool_calls must be a list"
|
||||
assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json"
|
||||
|
||||
|
||||
def test_tool_calls_from_generator_are_materialised():
|
||||
"""tool_calls passed as a generator must be converted to list on validation."""
|
||||
tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}')
|
||||
|
||||
def tool_calls_gen():
|
||||
yield tool_call
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "Search for vllm"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": tool_calls_gen(), # one-shot generator
|
||||
},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
|
||||
# Iterate twice — must not raise or return empty on second pass
|
||||
tool_calls_first = list(assistant_msg.get("tool_calls", []))
|
||||
tool_calls_second = list(assistant_msg.get("tool_calls", []))
|
||||
|
||||
assert len(tool_calls_first) == 1, "First read must return the tool call"
|
||||
assert len(tool_calls_second) == 1, "Second read must also return the tool call"
|
||||
|
||||
|
||||
def test_tool_calls_list_passthrough():
|
||||
"""tool_calls already provided as a list must remain a list."""
|
||||
tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}')
|
||||
messages = [
|
||||
{"role": "user", "content": "Calculate 2+2"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
assert isinstance(assistant_msg.get("tool_calls"), list)
|
||||
|
||||
|
||||
def test_messages_without_tool_calls_unaffected():
|
||||
"""Messages without tool_calls must be handled correctly."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
# None of the messages should have tool_calls injected
|
||||
for msg in req.messages:
|
||||
assert isinstance(msg, dict)
|
||||
assert msg.get("tool_calls") is None or msg.get("tool_calls") == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tool_calls", [1, 3])
|
||||
def test_multiple_tool_calls_materialised(num_tool_calls: int):
|
||||
"""Multiple tool calls in a single message are all preserved."""
|
||||
tool_calls = [
|
||||
_make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}')
|
||||
for i in range(num_tool_calls)
|
||||
]
|
||||
messages = [
|
||||
{"role": "user", "content": "Do things"},
|
||||
{"role": "assistant", "content": None, "tool_calls": iter(tool_calls)},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
|
||||
result_tool_calls = assistant_msg.get("tool_calls")
|
||||
assert isinstance(result_tool_calls, list)
|
||||
assert len(result_tool_calls) == num_tool_calls
|
||||
|
||||
# Verify after model_dump_json too
|
||||
_ = req.model_dump_json()
|
||||
assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls
|
||||
@@ -0,0 +1,158 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for multimodal features through the /inference/v1/generate endpoint.
|
||||
|
||||
Mirrors test_serving_tokens.py but exercises the multimodal piping
|
||||
using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
|
||||
/generate -> /detokenize path. Intentionally avoids running the HF
|
||||
processor in the pytest parent process to keep os.fork() in sibling
|
||||
tests (e.g. test_weight_transfer_llm.py) deadlock-free.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from PIL import Image
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.multimodal.utils import encode_image_url
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
|
||||
GEN_ENDPOINT = "/inference/v1/generate"
|
||||
RENDER_ENDPOINT = "/v1/chat/completions/render"
|
||||
DETOKENIZE_ENDPOINT = "/detokenize"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def test_image():
|
||||
return Image.new("RGB", (224, 224), color=(255, 0, 0))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--enforce-eager",
|
||||
"--no-enable-prefix-caching",
|
||||
]
|
||||
|
||||
envs = os.environ.copy()
|
||||
envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server: RemoteOpenAIServer):
|
||||
transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
|
||||
headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url=server.url_root,
|
||||
timeout=600,
|
||||
headers=headers,
|
||||
) as c:
|
||||
yield c
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_render_to_generate_roundtrip(client, test_image):
|
||||
"""End-to-end: render a multimodal chat -> feed into generate -> decode.
|
||||
|
||||
All preprocessing and detokenization happens in the server subprocess;
|
||||
the pytest parent never imports transformers or touches torch tensors.
|
||||
"""
|
||||
data_url = encode_image_url(test_image, format="PNG")
|
||||
|
||||
render_payload = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What color is this image? Answer in one word.",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
|
||||
render_resp.raise_for_status()
|
||||
render_data = render_resp.json()
|
||||
|
||||
# Validate render output structure: keys exist and values are non-empty
|
||||
# and well-typed.
|
||||
assert "token_ids" in render_data
|
||||
assert isinstance(render_data["token_ids"], list)
|
||||
assert len(render_data["token_ids"]) > 0
|
||||
assert all(isinstance(t, int) for t in render_data["token_ids"])
|
||||
|
||||
assert "features" in render_data
|
||||
features = render_data["features"]
|
||||
assert features is not None
|
||||
assert isinstance(features, dict)
|
||||
|
||||
assert "mm_hashes" in features
|
||||
assert "image" in features["mm_hashes"]
|
||||
image_hashes = features["mm_hashes"]["image"]
|
||||
assert isinstance(image_hashes, list)
|
||||
assert len(image_hashes) > 0
|
||||
assert all(isinstance(h, str) and h for h in image_hashes)
|
||||
|
||||
assert "mm_placeholders" in features
|
||||
assert "image" in features["mm_placeholders"]
|
||||
image_placeholders = features["mm_placeholders"]["image"]
|
||||
assert isinstance(image_placeholders, list)
|
||||
assert len(image_placeholders) > 0
|
||||
for p in image_placeholders:
|
||||
assert isinstance(p.get("offset"), int)
|
||||
assert isinstance(p.get("length"), int)
|
||||
assert p["length"] > 0
|
||||
|
||||
assert "kwargs_data" in features
|
||||
assert "image" in features["kwargs_data"]
|
||||
assert len(features["kwargs_data"]["image"]) > 0
|
||||
|
||||
# Build generate request from render output
|
||||
generate_payload = render_data
|
||||
generate_payload["sampling_params"] = {
|
||||
"max_tokens": 10,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
|
||||
gen_resp.raise_for_status()
|
||||
gen_data = gen_resp.json()
|
||||
|
||||
assert "choices" in gen_data
|
||||
assert isinstance(gen_data["choices"], list)
|
||||
assert len(gen_data["choices"]) >= 1
|
||||
choice = gen_data["choices"][0]
|
||||
assert "token_ids" in choice
|
||||
assert isinstance(choice["token_ids"], list)
|
||||
assert len(choice["token_ids"]) > 0
|
||||
assert all(isinstance(t, int) for t in choice["token_ids"])
|
||||
|
||||
detok_resp = await client.post(
|
||||
DETOKENIZE_ENDPOINT,
|
||||
json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
|
||||
)
|
||||
detok_resp.raise_for_status()
|
||||
detok_data = detok_resp.json()
|
||||
assert "prompt" in detok_data
|
||||
text = detok_data["prompt"]
|
||||
assert isinstance(text, str)
|
||||
assert len(text) > 0
|
||||
assert "red" in text.lower(), (
|
||||
f"Expected model to identify the red image, got: {text!r}"
|
||||
)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user