forked from Karylab-cklius/vllm
Compare commits
252
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
931b3f6110 | ||
|
|
95995bbef8 | ||
|
|
07351e0883 | ||
|
|
428b988c98 | ||
|
|
e54894fc85 | ||
|
|
bc2ae5a3d6 | ||
|
|
a474da2813 | ||
|
|
ce6a199ecc | ||
|
|
f88763efc3 | ||
|
|
333529deae | ||
|
|
8825608205 | ||
|
|
095d2f87e8 | ||
|
|
21792520e7 | ||
|
|
5e11b40365 | ||
|
|
f768b4473e | ||
|
|
914d0464c1 | ||
|
|
9f771b3ab9 | ||
|
|
c9d3c6e6af | ||
|
|
51adca74e6 | ||
|
|
e8eb0490ce | ||
|
|
e8ee2a78db | ||
|
|
2ec18f5df4 | ||
|
|
6dec49f27e | ||
|
|
b5587e1013 | ||
|
|
9ad5abe772 | ||
|
|
7d3195ea9f | ||
|
|
512f522192 | ||
|
|
4c34b2f6fc | ||
|
|
cf8a613a87 | ||
|
|
01acf96c6f | ||
|
|
079a4cf399 | ||
|
|
9744b699ba | ||
|
|
c662b4359e | ||
|
|
100c7b65e7 | ||
|
|
56bdf85e10 | ||
|
|
eba73068ea | ||
|
|
e9f331d72e | ||
|
|
c9bf77df92 | ||
|
|
3041344287 | ||
|
|
92762edc53 | ||
|
|
626daa2076 | ||
|
|
fe85a92e86 | ||
|
|
62b1bbe470 | ||
|
|
fa4b70555b | ||
|
|
447c372ac5 | ||
|
|
ff2c2bd80a | ||
|
|
cde8d24710 | ||
|
|
4a6dd1c3cc | ||
|
|
7ff65b1900 | ||
|
|
7f95a66cbf | ||
|
|
1b1c01de39 | ||
|
|
e9ba519f45 | ||
|
|
5ef33ab250 | ||
|
|
1c2c1eb8b9 | ||
|
|
8824f50f1f | ||
|
|
0098db9ec1 | ||
|
|
53ecc807c0 | ||
|
|
b7a2605020 | ||
|
|
d0009ddb0b | ||
|
|
424033f4fc | ||
|
|
da1e7311ca | ||
|
|
01cb41dcf5 | ||
|
|
2f314bc5e6 | ||
|
|
2196bac135 | ||
|
|
4b7869d6bc | ||
|
|
4a79262e0f | ||
|
|
3ed5231c6a | ||
|
|
9c2492e501 | ||
|
|
fe57be7809 | ||
|
|
8317cedc77 | ||
|
|
98a242ff61 | ||
|
|
e4ee48da2d | ||
|
|
342c58bc54 | ||
|
|
fe9c3d6c5f | ||
|
|
ccaf5ffaa3 | ||
|
|
0283f303d8 | ||
|
|
ac58e2a170 | ||
|
|
b8401a9bf4 | ||
|
|
9c271f9403 | ||
|
|
22fa63cfe8 | ||
|
|
8f87eb4622 | ||
|
|
cfa49213d7 | ||
|
|
29f64c5f5e | ||
|
|
eb6661d522 | ||
|
|
d622e27d2b | ||
|
|
5f76b3fb30 | ||
|
|
809d83c2dc | ||
|
|
33ef1941e2 | ||
|
|
a4905133f3 | ||
|
|
ecbe42e991 | ||
|
|
a250f1bd5f | ||
|
|
04eac6ba24 | ||
|
|
9047288b68 | ||
|
|
ed6d30377d | ||
|
|
6aa057c9d7 | ||
|
|
a2bd09c960 | ||
|
|
123674879e | ||
|
|
4254aeb56f | ||
|
|
aad88f8486 | ||
|
|
0210024ae7 | ||
|
|
4eafc72928 | ||
|
|
6d09769700 | ||
|
|
4506319a28 | ||
|
|
9b60e2ffaa | ||
|
|
3951d3eacd | ||
|
|
6f2c71be8f | ||
|
|
2463f00fb6 | ||
|
|
f946659fff | ||
|
|
f90aa44662 | ||
|
|
cefa5281a7 | ||
|
|
46794958f0 | ||
|
|
6ff8dea075 | ||
|
|
583e6f2226 | ||
|
|
96a85c5750 | ||
|
|
9db4650e5e | ||
|
|
5e584ce9ec | ||
|
|
1842447c09 | ||
|
|
16688b26a6 | ||
|
|
6fbec8ed47 | ||
|
|
5544f8c18b | ||
|
|
9f39b380d0 | ||
|
|
9a6a66f3b8 | ||
|
|
67eb6083e3 | ||
|
|
6ee081d1d0 | ||
|
|
66cc3fa559 | ||
|
|
6d85b36a9f | ||
|
|
ab5666eb7c | ||
|
|
f819265a4a | ||
|
|
936e0b79aa | ||
|
|
b2a5518679 | ||
|
|
908a713488 | ||
|
|
ec5ef0ac73 | ||
|
|
7b1e0b07d0 | ||
|
|
d249a9e90e | ||
|
|
d2e2e856ad | ||
|
|
766cb65d00 | ||
|
|
28c222157b | ||
|
|
3975eb6de6 | ||
|
|
5a94a19824 | ||
|
|
f95c11a848 | ||
|
|
257015d5e5 | ||
|
|
b47840019e | ||
|
|
989cc12d88 | ||
|
|
301024aa9c | ||
|
|
8256833fe6 | ||
|
|
8097591286 | ||
|
|
20d3743491 | ||
|
|
18563f2072 | ||
|
|
0e884fe638 | ||
|
|
fe5c115ee4 | ||
|
|
6867bcd076 | ||
|
|
c075702eae | ||
|
|
21b086d0aa | ||
|
|
3173441b0f | ||
|
|
8b1f3bebca | ||
|
|
2390caf157 | ||
|
|
87805fa11e | ||
|
|
304d5ba1a0 | ||
|
|
81d954f454 | ||
|
|
47fcb8ca68 | ||
|
|
191e3fdaa1 | ||
|
|
b9cf629bd0 | ||
|
|
3461c8b027 | ||
|
|
726efe177b | ||
|
|
595562651a | ||
|
|
3a30eaa1d7 | ||
|
|
fb5635d3f9 | ||
|
|
b42e878ec0 | ||
|
|
7243e02aa1 | ||
|
|
def8f52200 | ||
|
|
38fa87caca | ||
|
|
a023edfa5b | ||
|
|
b82fc1364d | ||
|
|
e06de7f005 | ||
|
|
cc3993b05d | ||
|
|
50dd4cb427 | ||
|
|
f774ba028a | ||
|
|
2aab9acf48 | ||
|
|
58631d7c3f | ||
|
|
a943839e9a | ||
|
|
6d8b80802b | ||
|
|
77fd2c8631 | ||
|
|
e729cc823d | ||
|
|
ec7aafc02a | ||
|
|
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 |
@@ -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
|
||||
|
||||
@@ -46,7 +46,7 @@ echo "Image not found, proceeding with build..."
|
||||
|
||||
# --- CUDA 13.0 for nightly builds ---
|
||||
# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
|
||||
NIGHTLY_CUDA_VERSION="13.0.0"
|
||||
NIGHTLY_CUDA_VERSION="13.0.2"
|
||||
NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
|
||||
NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
|
||||
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
group: Engine Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: Engine (1 GPU)
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/engine/
|
||||
- tests/v1/engine/
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
|
||||
@@ -0,0 +1,21 @@
|
||||
group: Kernels Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: vLLM IR Tests
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/ir
|
||||
- vllm/kernels
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s kernels/ir'
|
||||
@@ -0,0 +1,130 @@
|
||||
group: LoRA Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: LoRA Runtime + Utils
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_layers.py &&
|
||||
pytest -v -s lora/test_lora_checkpoints.py &&
|
||||
(pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
|
||||
pytest -v -s lora/test_lora_huggingface.py &&
|
||||
pytest -v -s lora/test_lora_manager.py &&
|
||||
pytest -v -s lora/test_lora_utils.py &&
|
||||
pytest -v -s lora/test_peft_helper.py &&
|
||||
pytest -v -s lora/test_resolver.py &&
|
||||
pytest -v -s lora/test_utils.py &&
|
||||
(pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
|
||||
(pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
|
||||
|
||||
- label: LoRA Fused/MoE Kernels
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_fused_moe_lora_kernel.py &&
|
||||
pytest -v -s lora/test_moe_lora_align_sum.py'
|
||||
|
||||
- label: LoRA Punica Kernels
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
set -o pipefail &&
|
||||
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'
|
||||
|
||||
- label: LoRA Punica FP8/XPU Ops
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_punica_ops_fp8.py &&
|
||||
pytest -v -s lora/test_punica_xpu_ops.py'
|
||||
|
||||
- label: LoRA Models
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
(pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
|
||||
pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
|
||||
pytest -v -s lora/test_qwen35_densemodel_lora.py &&
|
||||
pytest -v -s lora/test_transformers_model.py'
|
||||
|
||||
- label: LoRA Multimodal
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_default_mm_loras.py &&
|
||||
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
|
||||
(pytest -v -s lora/test_qwenvl.py || true) &&
|
||||
pytest -v -s lora/test_whisper.py'
|
||||
@@ -0,0 +1,55 @@
|
||||
group: Miscellaneous Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: V1 Core + KV + Metrics
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1/core
|
||||
- tests/v1/executor
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/metrics
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install -r requirements/kv_connectors.txt &&
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/executor'
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1/sample
|
||||
- tests/v1/logits_processors
|
||||
- tests/v1/test_oracle.py
|
||||
- tests/v1/test_request.py
|
||||
- tests/v1/test_outputs.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/logits_processors &&
|
||||
pytest -v -s v1/test_oracle.py &&
|
||||
pytest -v -s v1/test_request.py &&
|
||||
pytest -v -s v1/test_outputs.py'
|
||||
@@ -1,3 +1,13 @@
|
||||
# CUDA architecture lists — following PyTorch RELEASE.md
|
||||
# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
|
||||
# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
|
||||
env:
|
||||
CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
|
||||
# aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
|
||||
CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
|
||||
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
|
||||
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
|
||||
|
||||
steps:
|
||||
- input: "Provide Release version here"
|
||||
id: input-release-version
|
||||
@@ -14,12 +24,10 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -29,9 +37,7 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
|
||||
@@ -57,7 +63,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
@@ -70,7 +76,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
|
||||
@@ -108,96 +114,95 @@ steps:
|
||||
depends_on: block-build-release-images
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
- label: "Build release image - x86_64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0"
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0
|
||||
id: build-release-image-x86-cuda-12-9
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0
|
||||
id: build-release-image-arm64-cuda-12-9
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-ubuntu2404
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-ubuntu2404
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
id: build-release-image-x86-cuda-12-9-ubuntu2404
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
id: build-release-image-arm64-cuda-12-9-ubuntu2404
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
key: block-cpu-release-image-build
|
||||
@@ -238,7 +243,7 @@ steps:
|
||||
- group: "Publish release images"
|
||||
key: "publish-release-images"
|
||||
steps:
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
- label: "Create multi-arch manifest - CUDA 13.0"
|
||||
depends_on:
|
||||
- build-release-image-x86
|
||||
- build-release-image-arm64
|
||||
@@ -250,7 +255,7 @@ steps:
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
|
||||
- label: "Annotate release workflow - CUDA 12.9"
|
||||
- label: "Annotate release workflow - CUDA 13.0"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
id: annotate-release-workflow
|
||||
@@ -259,19 +264,19 @@ steps:
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-release.sh"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0"
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-13-0
|
||||
- build-release-image-arm64-cuda-13-0
|
||||
id: create-multi-arch-manifest-cuda-13-0
|
||||
- build-release-image-x86-cuda-12-9
|
||||
- build-release-image-arm64-cuda-12-9
|
||||
id: create-multi-arch-manifest-cuda-12-9
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-ubuntu2404
|
||||
- build-release-image-arm64-ubuntu2404
|
||||
@@ -283,17 +288,17 @@ steps:
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
- build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
|
||||
- build-release-image-x86-cuda-12-9-ubuntu2404
|
||||
- build-release-image-arm64-cuda-12-9-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-12-9-ubuntu2404
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub"
|
||||
depends_on:
|
||||
@@ -313,16 +318,16 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
|
||||
- label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest-cuda-13-0
|
||||
- create-multi-arch-manifest-cuda-12-9
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu129"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
|
||||
@@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
|
||||
(Optional) For CUDA 13.0:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
(Optional) For CUDA 12.9:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
|
||||
(Optional) For CPU:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
@@ -33,8 +33,8 @@ To download and upload the image:
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
|
||||
@@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
|
||||
@@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
|
||||
## ROCm
|
||||
|
||||
@@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu130
|
||||
docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker manifest push vllm/vllm-openai:latest-cu130
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
|
||||
docker manifest rm vllm/vllm-openai:latest-cu129
|
||||
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker manifest push vllm/vllm-openai:latest-cu129
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
|
||||
|
||||
docker manifest rm vllm/vllm-openai-cpu:latest || true
|
||||
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
|
||||
|
||||
@@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
|
||||
TORCH_INDEX_URL=""
|
||||
fi
|
||||
else
|
||||
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
|
||||
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
|
||||
fi
|
||||
echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ set -ex
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
INDICES_OUTPUT_DIR="indices"
|
||||
DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
|
||||
DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
|
||||
PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -25,22 +25,100 @@ export PYTHONPATH=".."
|
||||
###############################################################################
|
||||
|
||||
cleanup_docker() {
|
||||
# Share the same lock with image pull to avoid cleanup/pull races on one node.
|
||||
local docker_lock="/tmp/docker-pull.lock"
|
||||
exec 9>"$docker_lock"
|
||||
flock 9
|
||||
|
||||
docker_root=$(docker info -f '{{.DockerRootDir}}')
|
||||
if [ -z "$docker_root" ]; then
|
||||
echo "Failed to determine Docker root directory." >&2
|
||||
exit 1
|
||||
flock -u 9
|
||||
return 1
|
||||
fi
|
||||
echo "Docker root directory: $docker_root"
|
||||
|
||||
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
|
||||
threshold=70
|
||||
if [ "$disk_usage" -gt "$threshold" ]; then
|
||||
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
|
||||
docker image prune -f
|
||||
docker volume prune -f && docker system prune --force --filter "until=72h" --all
|
||||
echo "Docker images and volumes cleanup completed."
|
||||
echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
|
||||
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
|
||||
else
|
||||
echo "Disk usage is below $threshold%. No cleanup needed."
|
||||
echo "Disk usage is below $threshold%. Checking old CI images anyway."
|
||||
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
|
||||
fi
|
||||
echo "Old CI image cleanup completed."
|
||||
|
||||
flock -u 9
|
||||
}
|
||||
|
||||
cleanup_old_ci_images() {
|
||||
local repo_prefix="$1"
|
||||
local current_image_ref="$2"
|
||||
local ttl_hours="$3"
|
||||
local aggressive_cleanup="$4"
|
||||
|
||||
if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
|
||||
echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
|
||||
echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
|
||||
ttl_hours=72
|
||||
fi
|
||||
|
||||
local now_epoch cutoff_epoch
|
||||
now_epoch=$(date +%s)
|
||||
cutoff_epoch=$((now_epoch - ttl_hours * 3600))
|
||||
|
||||
local -a used_image_ids
|
||||
mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
|
||||
|
||||
local removed_count=0
|
||||
local examined_count=0
|
||||
declare -A seen_ids=()
|
||||
|
||||
while read -r image_ref image_id; do
|
||||
[[ -z "$image_ref" || -z "$image_id" ]] && continue
|
||||
((examined_count++))
|
||||
|
||||
# Keep the image this job is going to use.
|
||||
if [[ "$image_ref" == "$current_image_ref" ]]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Avoid duplicate deletes when multiple tags point to same image id.
|
||||
if [[ -n "${seen_ids[$image_id]:-}" ]]; then
|
||||
continue
|
||||
fi
|
||||
seen_ids[$image_id]=1
|
||||
|
||||
# Never delete images that are used by any container on this node.
|
||||
if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
local created created_epoch
|
||||
created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
|
||||
[[ -z "$created" ]] && continue
|
||||
created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
|
||||
[[ -z "$created_epoch" ]] && continue
|
||||
|
||||
if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
|
||||
if docker image rm -f "$image_id" >/dev/null 2>&1; then
|
||||
((removed_count++))
|
||||
fi
|
||||
fi
|
||||
done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
|
||||
|
||||
# Also trim old dangling layers; this is safe and does not remove referenced images.
|
||||
docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
|
||||
|
||||
if [[ "$aggressive_cleanup" == "1" ]]; then
|
||||
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
|
||||
else
|
||||
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
|
||||
fi
|
||||
}
|
||||
|
||||
@@ -240,7 +318,6 @@ fi
|
||||
cleanup_docker
|
||||
|
||||
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
|
||||
|
||||
# --- Build or pull test image ---
|
||||
IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
|
||||
@@ -266,8 +343,6 @@ fi
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
docker system prune -f || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
@@ -283,6 +358,7 @@ docker run \
|
||||
--ipc=host \
|
||||
--privileged \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
|
||||
--entrypoint="" \
|
||||
-e "HF_TOKEN=${HF_TOKEN:-}" \
|
||||
-e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \
|
||||
|
||||
@@ -12,9 +12,7 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true;
|
||||
docker image rm -f "${image_name}" || true;
|
||||
docker system prune -f || true;
|
||||
docker rm -f "${container_name}" || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euxo pipefail
|
||||
|
||||
# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
|
||||
THRESHOLD=${1:-0.8}
|
||||
NUM_Q=${2:-1319}
|
||||
PORT=${3:-8050}
|
||||
OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
|
||||
mkdir -p "${OUT_DIR}"
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
timeout 600 bash -c '
|
||||
until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
|
||||
sleep 1
|
||||
done'
|
||||
}
|
||||
|
||||
MODEL="Qwen/Qwen3-30B-A3B-FP8"
|
||||
BACK="allgather_reducescatter"
|
||||
|
||||
cleanup() {
|
||||
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
|
||||
kill "${SERVER_PID}" 2>/dev/null || true
|
||||
for _ in {1..20}; do
|
||||
kill -0 "${SERVER_PID}" 2>/dev/null || break
|
||||
sleep 0.5
|
||||
done
|
||||
kill -9 "${SERVER_PID}" 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
VLLM_DEEP_GEMM_WARMUP=skip \
|
||||
vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--data-parallel-size 4 \
|
||||
--enable-expert-parallel \
|
||||
--enable-eplb \
|
||||
--all2all-backend "$BACK" \
|
||||
--eplb-config '{"window_size":20, "step_interval":100, "use_async":true}' \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
assert acc >= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}"
|
||||
PY
|
||||
+2301
-2654
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,98 @@
|
||||
group: Disaggregated
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
@@ -196,6 +196,8 @@ steps:
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(B200)
|
||||
device: b200
|
||||
@@ -224,91 +226,6 @@ steps:
|
||||
commands:
|
||||
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
|
||||
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -29,6 +29,15 @@ steps:
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
|
||||
|
||||
- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 4
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
|
||||
|
||||
- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -28,6 +28,7 @@ steps:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
|
||||
@@ -12,6 +12,17 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
|
||||
|
||||
- label: Spec Decode Eagle Nightly B200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
|
||||
|
||||
- label: Spec Decode Speculators + MTP
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
@@ -23,6 +34,18 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: Spec Decode Speculators + MTP Nightly B200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/transformers_utils/configs/speculators/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: Spec Decode Ngram + Suffix
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
@@ -43,6 +66,17 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: Spec Decode Draft Model Nightly B200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- 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
|
||||
|
||||
+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
|
||||
|
||||
+1
-1
@@ -262,7 +262,7 @@ pull_request_rules:
|
||||
- files~=^docker/Dockerfile.xpu
|
||||
- files~=^\\.buildkite/intel_jobs/
|
||||
- files=\.buildkite/ci_config_intel.yaml
|
||||
- files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
|
||||
- files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
|
||||
- files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
|
||||
- files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
|
||||
- files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
|
||||
|
||||
@@ -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)`);
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
|
||||
# Limit the number of parallel jobs to avoid OOM
|
||||
export MAX_JOBS=1
|
||||
# Make sure release wheels are built for the following architectures
|
||||
export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
|
||||
export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
|
||||
|
||||
bash tools/check_repo.sh
|
||||
|
||||
|
||||
+25
-6
@@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
|
||||
# Supported python versions. These versions will be searched in order, the
|
||||
# first match will be selected. These should be kept in sync with setup.py.
|
||||
#
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
|
||||
|
||||
# Supported AMD GPU architectures.
|
||||
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
|
||||
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
|
||||
|
||||
# ROCm installation prefix. Default to /opt/rocm but allow override via
|
||||
# -DROCM_PATH=/your/rocm/path when invoking cmake.
|
||||
@@ -94,12 +94,15 @@ find_package(Torch REQUIRED)
|
||||
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
|
||||
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
|
||||
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
|
||||
# to support the whole generation without specifying all sub-architectures
|
||||
# see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")
|
||||
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1")
|
||||
else()
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
endif()
|
||||
|
||||
#
|
||||
@@ -923,6 +926,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 +955,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}")
|
||||
|
||||
@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
|
||||
| <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
|
||||
</p>
|
||||
|
||||
🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
|
||||
🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
|
||||
For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
|
||||
|
||||
---
|
||||
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
|
||||
- Efficient multi-LoRA support for dense and MoE layers
|
||||
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
||||
|
||||
vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
|
||||
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
||||
|
||||
- Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
|
||||
- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
|
||||
|
||||
@@ -404,6 +404,7 @@ def _build_attention_metadata(
|
||||
query_start_loc=q_start_gpu,
|
||||
query_start_loc_cpu=q_start_cpu,
|
||||
seq_lens=seq_lens_gpu,
|
||||
seq_lens_cpu_upper_bound=seq_lens_cpu,
|
||||
_seq_lens_cpu=seq_lens_cpu,
|
||||
_num_computed_tokens_cpu=num_computed_tokens_cpu,
|
||||
slot_mapping=slot_mapping,
|
||||
|
||||
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
fp8_w8a8_moe_quant_config,
|
||||
nvfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
|
||||
CutlassExpertsFp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
|
||||
|
||||
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import (
|
||||
fused_experts,
|
||||
fused_topk,
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Generic benchmark harness for vLLM IR ops.
|
||||
|
||||
Usage:
|
||||
python benchmarks/kernels/ir/bench_ir_ops.py
|
||||
python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm
|
||||
python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul
|
||||
python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph
|
||||
python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import csv
|
||||
import dataclasses
|
||||
import datetime
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
# Ensure repo root is on sys.path so `benchmarks` is importable as a package.
|
||||
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
|
||||
if _REPO_ROOT not in sys.path:
|
||||
sys.path.insert(0, _REPO_ROOT)
|
||||
|
||||
# Suppress noisy C++ warnings from vllm kernel registration (written to fd 2
|
||||
# directly by the dynamic linker, so Python-level sys.stderr redirect won't
|
||||
# catch them).
|
||||
_saved_fd = os.dup(2)
|
||||
try:
|
||||
with open(os.devnull, "w") as _devnull:
|
||||
os.dup2(_devnull.fileno(), 2)
|
||||
import torch
|
||||
|
||||
import vllm.kernels # noqa: E402, F401
|
||||
finally:
|
||||
os.dup2(_saved_fd, 2)
|
||||
os.close(_saved_fd)
|
||||
|
||||
from tqdm import tqdm # noqa: E402
|
||||
|
||||
from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS # noqa: E402 # isort: skip
|
||||
from vllm.ir.op import IrOp # noqa: E402
|
||||
from vllm.platforms import current_platform # noqa: E402
|
||||
from vllm.triton_utils import triton # noqa: E402
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class BenchConfig:
|
||||
use_cuda_graph: bool = True
|
||||
warmup: int = 25
|
||||
rep: int = 100
|
||||
|
||||
|
||||
def _pkg_version(name: str) -> str:
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
|
||||
with contextlib.suppress(PackageNotFoundError):
|
||||
return version(name)
|
||||
return "not installed"
|
||||
|
||||
|
||||
_METADATA_LABELS = {
|
||||
"timestamp": "Timestamp",
|
||||
"git_commit": "Git commit",
|
||||
"vllm": "vLLM",
|
||||
"pytorch": "PyTorch",
|
||||
"cuda_runtime": "CUDA runtime",
|
||||
"triton": "Triton",
|
||||
"cutlass": "CUTLASS",
|
||||
"helion": "Helion",
|
||||
"device": "Device",
|
||||
"bench_mode": "Bench mode",
|
||||
"warmup": "Warmup",
|
||||
"rep": "Repetitions",
|
||||
}
|
||||
|
||||
|
||||
def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]:
|
||||
from vllm.collect_env import get_env_info
|
||||
|
||||
env = get_env_info()
|
||||
|
||||
git_sha = "unknown"
|
||||
with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError):
|
||||
git_sha = (
|
||||
subprocess.check_output(
|
||||
["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
|
||||
device_name = current_platform.get_device_name()
|
||||
|
||||
warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)"
|
||||
rep_note = " replays" if cfg.use_cuda_graph else " ms"
|
||||
|
||||
return {
|
||||
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"git_commit": git_sha,
|
||||
"vllm": str(env.vllm_version),
|
||||
"pytorch": str(env.torch_version),
|
||||
"cuda_runtime": str(env.cuda_runtime_version),
|
||||
"triton": triton.__version__,
|
||||
"cutlass": _pkg_version("nvidia-cutlass-dsl"),
|
||||
"helion": _pkg_version("helion"),
|
||||
"device": device_name,
|
||||
"bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager",
|
||||
"warmup": f"{cfg.warmup}{warmup_note}",
|
||||
"rep": f"{cfg.rep}{rep_note}",
|
||||
}
|
||||
|
||||
|
||||
def print_metadata(metadata: dict[str, str]):
|
||||
print("=" * 60)
|
||||
for key, val in metadata.items():
|
||||
print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
def _clone_args(args: tuple) -> tuple:
|
||||
return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
|
||||
|
||||
|
||||
# TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as
|
||||
# inplace=True will mutate bench_args across iterations. Both CUDA graph
|
||||
# and eager modes will accumulate drift from repeated in-place mutation.
|
||||
# We need to re-clone inputs per iteration for inplace ops.
|
||||
def _bench_one(fn, args, cfg: BenchConfig) -> float:
|
||||
bench_args = _clone_args(args)
|
||||
bench_fn = lambda: fn(*bench_args)
|
||||
|
||||
if cfg.use_cuda_graph:
|
||||
ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5])
|
||||
else:
|
||||
ms = triton.testing.do_bench(
|
||||
bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5]
|
||||
)
|
||||
return ms * 1000
|
||||
|
||||
|
||||
# TODO(gmagogsfm): Once compiled native implementation lands (#38775),
|
||||
# the benchmark baseline should be the compiled native (what vLLM runs by
|
||||
# default) rather than the uncompiled native implementation.
|
||||
def collect_timings(
|
||||
op: IrOp, shape_configs: list[dict], cfg: BenchConfig
|
||||
) -> tuple[list[str], list[str], dict[str, dict[str, float]]]:
|
||||
def fmt(v) -> str:
|
||||
return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v)
|
||||
|
||||
case_names = [
|
||||
"_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs
|
||||
]
|
||||
providers = [n for n, impl in op.impls.items() if impl.supported]
|
||||
|
||||
results: dict[str, dict[str, float]] = {c: {} for c in case_names}
|
||||
for provider in providers:
|
||||
impl = op.impls[provider]
|
||||
desc = f"{op.name} / {provider}"
|
||||
for case_name, kwargs in tqdm(
|
||||
zip(case_names, shape_configs),
|
||||
desc=desc,
|
||||
total=len(case_names),
|
||||
unit=" cases",
|
||||
):
|
||||
args = op.generate_inputs(**kwargs)
|
||||
if impl.supports_args(*args):
|
||||
results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg)
|
||||
else:
|
||||
results[case_name][provider] = float("nan")
|
||||
|
||||
return case_names, providers, results
|
||||
|
||||
|
||||
def analyze_results(
|
||||
op_name: str,
|
||||
case_names: list[str],
|
||||
providers: list[str],
|
||||
results: dict[str, dict[str, float]],
|
||||
) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]:
|
||||
native_col = "native"
|
||||
non_native = [p for p in providers if p != native_col]
|
||||
|
||||
header_cols = ["case"]
|
||||
for p in providers:
|
||||
header_cols.append(f"{p} (us)")
|
||||
for p in non_native:
|
||||
header_cols.append(f"{p} speedup")
|
||||
|
||||
detail_rows: list[dict[str, str]] = []
|
||||
speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native}
|
||||
|
||||
for case_name in case_names:
|
||||
timings = results[case_name]
|
||||
row: dict[str, str] = {"case": case_name}
|
||||
|
||||
for p in providers:
|
||||
val = timings.get(p, float("nan"))
|
||||
row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a"
|
||||
|
||||
native_us = timings.get(native_col, float("nan"))
|
||||
for p in non_native:
|
||||
p_us = timings.get(p, float("nan"))
|
||||
if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0:
|
||||
speedup = native_us / p_us
|
||||
row[f"{p} speedup"] = f"{speedup:.2f}x"
|
||||
speedup_data[p].append((speedup, case_name))
|
||||
else:
|
||||
row[f"{p} speedup"] = "n/a"
|
||||
|
||||
detail_rows.append(row)
|
||||
|
||||
summary_rows: list[dict[str, str]] = []
|
||||
for p in non_native:
|
||||
entries = speedup_data[p]
|
||||
if not entries:
|
||||
continue
|
||||
speedups = [s for s, _ in entries]
|
||||
geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups))
|
||||
best_val, best_case = max(entries)
|
||||
worst_val, worst_case = min(entries)
|
||||
wins = sum(1 for s in speedups if s > 1.0)
|
||||
losses = sum(1 for s in speedups if s < 1.0)
|
||||
total = len(speedups)
|
||||
|
||||
print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):")
|
||||
print(f" geomean speedup: {geomean:.2f}x")
|
||||
print(f" best: {best_val:.2f}x ({best_case})")
|
||||
print(f" worst: {worst_val:.2f}x ({worst_case})")
|
||||
|
||||
summary_rows.append(
|
||||
{
|
||||
"op": op_name,
|
||||
"provider": p,
|
||||
"geomean_speedup": f"{geomean:.2f}",
|
||||
"best_speedup": f"{best_val:.2f}",
|
||||
"best_case": best_case,
|
||||
"worst_speedup": f"{worst_val:.2f}",
|
||||
"worst_case": worst_case,
|
||||
"wins": str(wins),
|
||||
"losses": str(losses),
|
||||
"total": str(total),
|
||||
}
|
||||
)
|
||||
|
||||
return detail_rows, summary_rows, header_cols
|
||||
|
||||
|
||||
def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]):
|
||||
with open(path, "w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def save_results(
|
||||
save_dir: str,
|
||||
op_name: str,
|
||||
detail_rows: list[dict[str, str]],
|
||||
header_cols: list[str],
|
||||
all_summary_rows: list[dict[str, str]],
|
||||
metadata: dict[str, str],
|
||||
):
|
||||
write_csv(
|
||||
os.path.join(save_dir, f"{op_name}_detail.csv"),
|
||||
detail_rows,
|
||||
header_cols,
|
||||
)
|
||||
if all_summary_rows:
|
||||
write_csv(
|
||||
os.path.join(save_dir, "summary.csv"),
|
||||
all_summary_rows,
|
||||
list(all_summary_rows[0].keys()),
|
||||
)
|
||||
write_csv(
|
||||
os.path.join(save_dir, "metadata.csv"),
|
||||
[metadata],
|
||||
list(metadata.keys()),
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops")
|
||||
parser.add_argument(
|
||||
"--ops",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Comma-separated list of op names to benchmark (substring match)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-cuda-graph",
|
||||
action="store_true",
|
||||
help="Disable CUDA graph; use do_bench with L2 cache flushing instead",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=25,
|
||||
help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rep",
|
||||
type=int,
|
||||
default=100,
|
||||
help="Repetition time in ms (do_bench) or number of graph replays "
|
||||
"(do_bench_cudagraph) (default: 100)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory to save results (default: auto-created temp dir)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
cfg = BenchConfig(
|
||||
use_cuda_graph=not args.no_cuda_graph,
|
||||
warmup=args.warmup,
|
||||
rep=args.rep,
|
||||
)
|
||||
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
|
||||
metadata = collect_env_metadata(cfg)
|
||||
print_metadata(metadata)
|
||||
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
save_dir = args.save_path or os.path.join(
|
||||
tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}"
|
||||
)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None
|
||||
all_summary_rows: list[dict[str, str]] = []
|
||||
|
||||
for op in IrOp.registry.values():
|
||||
if op_filters and not any(f in op.name for f in op_filters):
|
||||
continue
|
||||
if not op.has_input_generator:
|
||||
print(f"Skipping op '{op.name}': no input generator registered")
|
||||
continue
|
||||
if op.name not in SHAPE_CONFIGS:
|
||||
raise RuntimeError(
|
||||
f"No benchmark shape config for op '{op.name}'. "
|
||||
f"Add it to benchmarks/kernels/ir/shapes.py"
|
||||
)
|
||||
|
||||
case_names, providers, results = collect_timings(
|
||||
op, SHAPE_CONFIGS[op.name], cfg
|
||||
)
|
||||
detail_rows, summary_rows, header_cols = analyze_results(
|
||||
op.name, case_names, providers, results
|
||||
)
|
||||
all_summary_rows.extend(summary_rows)
|
||||
|
||||
save_results(
|
||||
save_dir,
|
||||
op.name,
|
||||
detail_rows,
|
||||
header_cols,
|
||||
all_summary_rows,
|
||||
metadata,
|
||||
)
|
||||
|
||||
print(f"\nResults saved to: {save_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Shape configurations for IR op benchmarks.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384]
|
||||
COMMON_HIDDEN_SIZES = [
|
||||
2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
|
||||
3072, # Gemma 7B/9B
|
||||
4096, # Llama 3 8B, Qwen 3 8B, Mistral 7B
|
||||
5120, # Llama 4 Scout 17B-16E
|
||||
7168, # DeepSeek V3
|
||||
8192, # Llama 3 70B
|
||||
16384, # Llama 3 405B
|
||||
]
|
||||
|
||||
# Each entry maps an op name to a list of kwarg dicts that will be passed
|
||||
# to that op's registered input generator via op.generate_inputs(**kwargs).
|
||||
SHAPE_CONFIGS: dict[str, list[dict]] = {
|
||||
"rms_norm": [
|
||||
{"num_tokens": n, "hidden_size": d, "dtype": dtype}
|
||||
for dtype in [torch.float16, torch.bfloat16, torch.float32]
|
||||
for d in COMMON_HIDDEN_SIZES
|
||||
for n in NUM_TOKENS
|
||||
],
|
||||
}
|
||||
+56
-22
@@ -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)
|
||||
@@ -146,16 +161,49 @@ elseif (S390_FOUND)
|
||||
"-mtune=native")
|
||||
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
message(STATUS "RISC-V detected")
|
||||
if(RVV_BF16_FOUND)
|
||||
message(STATUS "BF16 extension detected")
|
||||
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
|
||||
add_compile_definitions(RISCV_BF16_SUPPORT)
|
||||
elseif (RVV_FP16_FOUND)
|
||||
message(WARNING "BF16 functionality is not available")
|
||||
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
|
||||
# VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
|
||||
# by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
|
||||
if(NOT DEFINED VLLM_RVV_VLEN)
|
||||
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
|
||||
if(EXISTS /proc/cpuinfo)
|
||||
file(READ /proc/cpuinfo _cpuinfo)
|
||||
set(_best 0)
|
||||
foreach(_n IN ITEMS 128 256 512 1024)
|
||||
if(_cpuinfo MATCHES "zvl${_n}b")
|
||||
set(_best ${_n})
|
||||
endif()
|
||||
endforeach()
|
||||
if(_best GREATER 0)
|
||||
set(VLLM_RVV_VLEN ${_best})
|
||||
endif()
|
||||
endif()
|
||||
# If auto-detect failed (no /proc/cpuinfo or no zvl<N>b reported)
|
||||
# but the compiler supports RVV, require explicit specification.
|
||||
if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
|
||||
message(FATAL_ERROR
|
||||
"RISC-V RVV is available but VLEN could not be auto-detected. "
|
||||
"Please specify VLEN explicitly:\n"
|
||||
" -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
|
||||
" -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n"
|
||||
" -DVLLM_RVV_VLEN=0 (force scalar, no RVV)")
|
||||
endif()
|
||||
endif()
|
||||
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
|
||||
message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
|
||||
if(RVV_BF16_FOUND)
|
||||
message(STATUS "BF16 extension detected")
|
||||
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
|
||||
add_compile_definitions(RISCV_BF16_SUPPORT)
|
||||
elseif(RVV_FP16_FOUND)
|
||||
message(WARNING "BF16 functionality is not available")
|
||||
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
|
||||
else()
|
||||
message(STATUS "compile riscv with scalar (no FP16/BF16)")
|
||||
set(MARCH_FLAGS -march=rv64gc)
|
||||
endif()
|
||||
else()
|
||||
message(STATUS "compile riscv with scalar")
|
||||
list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
|
||||
set(MARCH_FLAGS -march=rv64gc)
|
||||
endif()
|
||||
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
|
||||
else()
|
||||
@@ -175,20 +223,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}")
|
||||
|
||||
+40
-10
@@ -599,6 +599,11 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
|
||||
const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
|
||||
// Find batch index within a block
|
||||
__shared__ int batch_idx[BLOCK_Y_SIZE];
|
||||
if (threadIdx.x == 0) {
|
||||
batch_idx[threadIdx.y] = -1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
|
||||
iter++) {
|
||||
int tid = iter * blockDim.x + threadIdx.x;
|
||||
@@ -611,16 +616,18 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
__syncwarp();
|
||||
#endif
|
||||
__syncthreads();
|
||||
|
||||
if (head_idx >= head_dim || token_idx >= num_tokens) {
|
||||
// num_tokens may be an allocation upper bound when Python avoids a D2H sync.
|
||||
// Only tokens covered by the exact device-side cu_seq_lens are valid to
|
||||
// gather.
|
||||
const int batch = batch_idx[threadIdx.y];
|
||||
if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
|
||||
return;
|
||||
}
|
||||
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
|
||||
const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
|
||||
inbatch_seq_idx / cache_block_size];
|
||||
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
|
||||
const int block_idx =
|
||||
block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
|
||||
const int64_t src_block_offset = block_idx * block_stride;
|
||||
const int64_t cache_inblock_offset =
|
||||
(inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
|
||||
@@ -724,6 +731,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 +770,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);
|
||||
@@ -1470,6 +1497,9 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
|
||||
TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
|
||||
TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
|
||||
TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
|
||||
TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
|
||||
ql_nope.scalar_type() == at::ScalarType::BFloat16,
|
||||
"ql_nope must be float16 or bfloat16 dtype");
|
||||
|
||||
if (num_tokens == 0) return;
|
||||
|
||||
@@ -1481,7 +1511,7 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
|
||||
VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
|
||||
vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
|
||||
q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
|
||||
q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
|
||||
|
||||
@@ -61,8 +61,23 @@
|
||||
#endif
|
||||
|
||||
#ifdef __aarch64__
|
||||
// Implementation copied from Arm Optimized Routines (expf AdvSIMD)
|
||||
// Implementation of neon_expf copied from Arm Optimized Routines (expf
|
||||
// AdvSIMD)
|
||||
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
|
||||
//
|
||||
// Additional fast exponential intended for cases where outputs will be
|
||||
// downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
|
||||
// for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
|
||||
// clamps inputs outside this range to 0 / inf. Implementation is similar to
|
||||
// exp_u20, but:
|
||||
// - uses a third degree polynomial approximation for exp(r) instead of a
|
||||
// fifth degree one, with coefficients re-tuned.
|
||||
// - does not split natural log (ln) into high / low parts
|
||||
// - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
|
||||
// exp(x) = 2^n (exp(r))
|
||||
// r = x - n*ln2, with n = round(x/ln2)
|
||||
// exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
|
||||
// n = round(x / ln2), r = x - n*ln2
|
||||
#include <limits>
|
||||
#define DEFINE_FAST_EXP \
|
||||
const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
|
||||
@@ -106,7 +121,38 @@
|
||||
result.val[2] = neon_expf(vec.reg.val[2]); \
|
||||
result.val[3] = neon_expf(vec.reg.val[3]); \
|
||||
return vec_op::FP32Vec16(result); \
|
||||
};
|
||||
}; \
|
||||
const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \
|
||||
const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \
|
||||
constexpr float ln2 = 0x1.62e43p-1f; \
|
||||
constexpr float f_c2 = 0x1.5592ecp-3f; \
|
||||
const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \
|
||||
auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \
|
||||
always_inline)) { \
|
||||
const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \
|
||||
const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \
|
||||
float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
|
||||
float32x4_t r = vfmsq_n_f32(values, n, ln2); \
|
||||
uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
|
||||
float32x4_t r2 = vmulq_f32(r, r); \
|
||||
float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \
|
||||
float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \
|
||||
float32x4_t p = vfmaq_f32(s, q, r2); \
|
||||
float32x4_t y = \
|
||||
vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \
|
||||
y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \
|
||||
y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \
|
||||
return y; \
|
||||
}; \
|
||||
auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \
|
||||
__attribute__((always_inline)) { \
|
||||
float32x4x4_t result; \
|
||||
result.val[0] = neon_expf_f16(vec.reg.val[0]); \
|
||||
result.val[1] = neon_expf_f16(vec.reg.val[1]); \
|
||||
result.val[2] = neon_expf_f16(vec.reg.val[2]); \
|
||||
result.val[3] = neon_expf_f16(vec.reg.val[3]); \
|
||||
return vec_op::FP32Vec16(result); \
|
||||
};
|
||||
|
||||
#endif // __aarch64__
|
||||
|
||||
|
||||
@@ -1152,7 +1152,11 @@ class AttentionMainLoop {
|
||||
bool use_sink) {
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
DEFINE_FAST_EXP
|
||||
bool constexpr IsReducedPrecision =
|
||||
std::is_same_v<query_t, c10::BFloat16> ||
|
||||
std::is_same_v<query_t, c10::Half>;
|
||||
#endif
|
||||
|
||||
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
|
||||
static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
|
||||
|
||||
@@ -1201,8 +1205,17 @@ class AttentionMainLoop {
|
||||
vec = vec - max_vec;
|
||||
|
||||
// compute exp
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
vec = fast_exp(vec);
|
||||
|
||||
#if defined(DEFINE_FAST_EXP)
|
||||
#ifdef __aarch64__
|
||||
if constexpr (IsReducedPrecision) {
|
||||
vec = fast_exp_f16(vec);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
|
||||
prob_buffer_vec_t output_vec(vec);
|
||||
output_vec.save(curr_prob_buffer_iter);
|
||||
#else
|
||||
@@ -1258,7 +1271,11 @@ class AttentionMainLoop {
|
||||
int32_t kv_tile_token_num, float softcap_scale) {
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
DEFINE_FAST_EXP
|
||||
bool constexpr IsReducedPrecision =
|
||||
std::is_same_v<query_t, c10::BFloat16> ||
|
||||
std::is_same_v<query_t, c10::Half>;
|
||||
#endif
|
||||
|
||||
float inv_softcap_scale = 1.0 / softcap_scale;
|
||||
vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
|
||||
vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
|
||||
@@ -1272,8 +1289,15 @@ class AttentionMainLoop {
|
||||
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
|
||||
vec = vec * inv_softcap_scale_vec;
|
||||
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
vec = fast_exp(vec);
|
||||
#if defined(DEFINE_FAST_EXP)
|
||||
#ifdef __aarch64__
|
||||
if constexpr (IsReducedPrecision) {
|
||||
vec = fast_exp_f16(vec);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
vec_op::FP32Vec16 inv_vec = ones_vec / vec;
|
||||
vec = (vec - inv_vec) / (vec + inv_vec);
|
||||
#else
|
||||
|
||||
+16
-823
@@ -1,832 +1,25 @@
|
||||
#ifndef CPU_TYPES_RISCV_HPP
|
||||
#define CPU_TYPES_RISCV_HPP
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <limits>
|
||||
#include <riscv_vector.h>
|
||||
#include <torch/all.h>
|
||||
// RISC-V Vector (RVV) CPU type definitions for vLLM.
|
||||
//
|
||||
// Supports multiple VLENs via compile-time dispatch. The compiler defines
|
||||
// __riscv_v_min_vlen from the zvl<N>b extension in -march. The defs header
|
||||
// maps VLEN to the correct LMUL suffixes, and the impl header provides
|
||||
// VLEN-independent class implementations.
|
||||
//
|
||||
// To add support for a new VLEN, add the LMUL mapping in
|
||||
// cpu_types_riscv_defs.hpp (the impl header needs no changes).
|
||||
|
||||
// ============================================================================
|
||||
// Vector Register Type Definitions (VLEN=128 bits)
|
||||
// ============================================================================
|
||||
|
||||
typedef vfloat16m1_t fixed_vfloat16m1_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef vfloat16m2_t fixed_vfloat16m2_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
|
||||
typedef vfloat32m1_t fixed_vfloat32m1_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef vfloat32m2_t fixed_vfloat32m2_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef vfloat32m4_t fixed_vfloat32m4_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
typedef vfloat32m8_t fixed_vfloat32m8_t
|
||||
__attribute__((riscv_rvv_vector_bits(1024)));
|
||||
|
||||
typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
typedef vuint16m1_t fixed_vuint16m1_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef vuint16m2_t fixed_vuint16m2_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef vuint16m4_t fixed_vuint16m4_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
typedef vbfloat16m1_t fixed_vbfloat16m1_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef vbfloat16m2_t fixed_vbfloat16m2_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef vbfloat16m4_t fixed_vbfloat16m4_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
#ifndef __riscv_vector
|
||||
#error "cpu_types_riscv.hpp included in a non-RVV translation unit"
|
||||
#endif
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
|
||||
#else
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
|
||||
#ifndef __riscv_v_min_vlen
|
||||
#error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
|
||||
#endif
|
||||
|
||||
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
||||
#include "cpu_types_riscv_defs.hpp"
|
||||
#include "cpu_types_riscv_impl.hpp"
|
||||
|
||||
#define FORCE_INLINE __attribute__((always_inline)) inline
|
||||
|
||||
namespace {
|
||||
template <typename T, T... indexes, typename F>
|
||||
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
|
||||
(f(std::integral_constant<T, indexes>{}), ...);
|
||||
};
|
||||
} // namespace
|
||||
|
||||
template <typename T, T count, typename F,
|
||||
typename = std::enable_if_t<std::is_invocable_v<F, T>>>
|
||||
constexpr void unroll_loop(F&& f) {
|
||||
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct Vec {
|
||||
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
|
||||
};
|
||||
|
||||
struct FP32Vec8;
|
||||
struct FP32Vec16;
|
||||
|
||||
// ============================================================================
|
||||
// FP16 Implementation
|
||||
// ============================================================================
|
||||
|
||||
struct FP16Vec8 : public Vec<FP16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_vfloat16m1_t reg;
|
||||
|
||||
explicit FP16Vec8(const void* ptr)
|
||||
: reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
|
||||
VEC_ELEM_NUM)) {};
|
||||
|
||||
explicit FP16Vec8(const FP32Vec8&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
__riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
__riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(_Float16);
|
||||
__riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_vfloat16m2_t reg;
|
||||
|
||||
explicit FP16Vec16(const void* ptr)
|
||||
: reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
|
||||
VEC_ELEM_NUM)) {};
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16& vec);
|
||||
|
||||
void save(void* ptr) const {
|
||||
__riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
__riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(_Float16);
|
||||
__riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// BF16 Implementation
|
||||
// ============================================================================
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
|
||||
FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
|
||||
return __riscv_vreinterpret_v_bf16m1_u16m1(v);
|
||||
}
|
||||
FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
|
||||
return __riscv_vreinterpret_v_bf16m2_u16m2(v);
|
||||
}
|
||||
FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
|
||||
return __riscv_vreinterpret_v_bf16m4_u16m4(v);
|
||||
}
|
||||
|
||||
struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_vbfloat16m1_t reg;
|
||||
|
||||
explicit BF16Vec8(const void* ptr)
|
||||
: reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
|
||||
explicit BF16Vec8(const FP32Vec8&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
__riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
__riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
|
||||
elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
__riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
|
||||
bf16_to_u16(reg), VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_vbfloat16m2_t reg;
|
||||
|
||||
explicit BF16Vec16(const void* ptr)
|
||||
: reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
__riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
__riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
|
||||
elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
__riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
|
||||
bf16_to_u16(reg), VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
constexpr static int VEC_ELEM_NUM = 32;
|
||||
fixed_vbfloat16m4_t reg;
|
||||
|
||||
explicit BF16Vec32(const void* ptr)
|
||||
: reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
|
||||
fixed_vuint16m4_t u16_combined =
|
||||
__riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
|
||||
reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
|
||||
};
|
||||
|
||||
void save(void* ptr) const {
|
||||
__riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
__riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
|
||||
elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
__riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
|
||||
bf16_to_u16(reg), VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
#else
|
||||
// ============================================================================
|
||||
// BF16 Fallback Implementation (FP32 Simulation)
|
||||
// ============================================================================
|
||||
|
||||
struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_vfloat32m2_t reg_fp32;
|
||||
explicit BF16Vec8(const void* ptr) {
|
||||
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
|
||||
float tmp[8];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
|
||||
std::memcpy(&tmp[i], &v, 4);
|
||||
}
|
||||
reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
|
||||
}
|
||||
explicit BF16Vec8(const FP32Vec8&);
|
||||
void save(void* ptr) const {
|
||||
float tmp[8];
|
||||
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
float tmp[8];
|
||||
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < elem_num; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
float tmp[8];
|
||||
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
|
||||
uint8_t* u8 = static_cast<uint8_t*>(ptr);
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
uint16_t val = static_cast<uint16_t>(v >> 16);
|
||||
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_vfloat32m4_t reg_fp32;
|
||||
explicit BF16Vec16(const void* ptr) {
|
||||
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
|
||||
float tmp[16];
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
|
||||
std::memcpy(&tmp[i], &v, 4);
|
||||
}
|
||||
reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
|
||||
}
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
void save(void* ptr) const {
|
||||
float tmp[16];
|
||||
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
float tmp[16];
|
||||
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < elem_num; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
float tmp[16];
|
||||
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
|
||||
uint8_t* u8 = static_cast<uint8_t*>(ptr);
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
uint16_t val = static_cast<uint16_t>(v >> 16);
|
||||
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
constexpr static int VEC_ELEM_NUM = 32;
|
||||
fixed_vfloat32m8_t reg_fp32;
|
||||
|
||||
explicit BF16Vec32(const void* ptr) {
|
||||
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
|
||||
float tmp[32];
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
|
||||
std::memcpy(&tmp[i], &v, 4);
|
||||
}
|
||||
reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
|
||||
}
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
float tmp_small[8];
|
||||
__riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
|
||||
float tmp_large[32];
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
|
||||
}
|
||||
reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
|
||||
}
|
||||
|
||||
void save(void* ptr) const {
|
||||
float tmp[32];
|
||||
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
|
||||
void save(void* ptr, int elem_num) const {
|
||||
float tmp[32];
|
||||
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < elem_num; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
float tmp[32];
|
||||
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
|
||||
uint8_t* u8 = static_cast<uint8_t*>(ptr);
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
uint16_t val = static_cast<uint16_t>(v >> 16);
|
||||
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
|
||||
}
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
// ============================================================================
|
||||
// FP32 Implementation
|
||||
// ============================================================================
|
||||
|
||||
struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
constexpr static int VEC_ELEM_NUM = 4;
|
||||
fixed_vfloat32m1_t reg;
|
||||
explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec4(const float* ptr)
|
||||
: reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
|
||||
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
|
||||
void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
|
||||
void save(float* ptr, int elem_num) const {
|
||||
__riscv_vse32_v_f32m1(ptr, reg, elem_num);
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_vfloat32m2_t reg;
|
||||
|
||||
explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(const float* ptr)
|
||||
: reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
|
||||
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
|
||||
explicit FP32Vec8(const FP16Vec8& v)
|
||||
: reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(fixed_vfloat16m1_t v)
|
||||
: reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
explicit FP32Vec8(fixed_vbfloat16m1_t v)
|
||||
: reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(const BF16Vec8& v)
|
||||
: reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
|
||||
#else
|
||||
explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
|
||||
#endif
|
||||
|
||||
float reduce_sum() const {
|
||||
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
|
||||
scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
FP32Vec8 operator*(const FP32Vec8& b) const {
|
||||
return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 operator+(const FP32Vec8& b) const {
|
||||
return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 operator-(const FP32Vec8& b) const {
|
||||
return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 operator/(const FP32Vec8& b) const {
|
||||
return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 min(const FP32Vec8& b) const {
|
||||
return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 max(const FP32Vec8& b) const {
|
||||
return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 abs() const {
|
||||
return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
|
||||
return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
|
||||
}
|
||||
FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
|
||||
return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
|
||||
}
|
||||
|
||||
FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
|
||||
fixed_vfloat32m2_t temp =
|
||||
__riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
|
||||
return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
|
||||
void save(float* ptr, int elem_num) const {
|
||||
__riscv_vse32_v_f32m2(ptr, reg, elem_num);
|
||||
}
|
||||
void save_strided(float* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(float);
|
||||
__riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
|
||||
FP32Vec8 exp() const {
|
||||
const float inv_ln2 = 1.44269504088896341f;
|
||||
fixed_vfloat32m2_t x_scaled =
|
||||
__riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
|
||||
fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m2_t r =
|
||||
__riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m2_t poly =
|
||||
__riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vint32m2_t biased_exp =
|
||||
__riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
|
||||
biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
|
||||
fixed_vint32m2_t exponent_bits =
|
||||
__riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m2_t scale =
|
||||
__riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
|
||||
|
||||
return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 tanh() const {
|
||||
fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
|
||||
__riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m2_t x2 =
|
||||
__riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
|
||||
FP32Vec8 exp_val = FP32Vec8(x2).exp();
|
||||
fixed_vfloat32m2_t num =
|
||||
__riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m2_t den =
|
||||
__riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
|
||||
return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 er() const {
|
||||
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
|
||||
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
|
||||
fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
|
||||
__riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
|
||||
t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
|
||||
a4, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
|
||||
a3, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
|
||||
a2, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
|
||||
a1, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m2_t exp_val =
|
||||
FP32Vec8(__riscv_vfneg_v_f32m2(
|
||||
__riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM))
|
||||
.exp()
|
||||
.reg;
|
||||
fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
|
||||
__riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
|
||||
return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_vfloat32m4_t reg;
|
||||
|
||||
explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(const float* ptr)
|
||||
: reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
|
||||
explicit FP32Vec16(const FP32Vec8& data)
|
||||
: reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
|
||||
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
explicit FP32Vec16(fixed_vbfloat16m2_t v)
|
||||
: reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(const BF16Vec16& v)
|
||||
: reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
|
||||
#else
|
||||
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
|
||||
#endif
|
||||
|
||||
FP32Vec16 operator+(const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 operator-(const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 operator*(const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 operator/(const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
float reduce_sum() const {
|
||||
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
|
||||
scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
float reduce_max() const {
|
||||
fixed_vfloat32m1_t scalar =
|
||||
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
|
||||
scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
float reduce_min() const {
|
||||
fixed_vfloat32m1_t scalar =
|
||||
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
|
||||
scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
template <int group_size>
|
||||
float reduce_sub_sum(int idx) {
|
||||
static_assert(VEC_ELEM_NUM % group_size == 0);
|
||||
const int start = idx * group_size;
|
||||
vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
|
||||
vbool8_t mask = __riscv_vmand_mm_b8(
|
||||
__riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
|
||||
__riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM);
|
||||
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
|
||||
scalar =
|
||||
__riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
};
|
||||
|
||||
FP32Vec16 max(const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 min(const FP32Vec16& b) const {
|
||||
return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 abs() const {
|
||||
return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
|
||||
return FP32Vec16(__riscv_vfmin_vv_f32m4(
|
||||
max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
|
||||
void save(float* ptr, int elem_num) const {
|
||||
__riscv_vse32_v_f32m4(ptr, reg, elem_num);
|
||||
}
|
||||
void save_strided(float* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(float);
|
||||
__riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
|
||||
FP32Vec16 exp() const {
|
||||
const float inv_ln2 = 1.44269504088896341f;
|
||||
fixed_vfloat32m4_t x_scaled =
|
||||
__riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
|
||||
fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m4_t r =
|
||||
__riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m4_t poly =
|
||||
__riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
|
||||
0.009618129107628f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
|
||||
0.055504108664821f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
|
||||
0.240226506959101f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
|
||||
0.693147180559945f, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
|
||||
1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
|
||||
__riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
|
||||
__riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
|
||||
|
||||
return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 tanh() const {
|
||||
fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
|
||||
__riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
|
||||
FP32Vec16 exp_val =
|
||||
FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
|
||||
return FP32Vec16(__riscv_vfdiv_vv_f32m4(
|
||||
__riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
|
||||
__riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 er() const {
|
||||
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
|
||||
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
|
||||
fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
|
||||
fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
|
||||
__riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
|
||||
1.0f, VEC_ELEM_NUM),
|
||||
1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
|
||||
a4, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
|
||||
a3, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
|
||||
a2, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
|
||||
a1, VEC_ELEM_NUM);
|
||||
poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
|
||||
|
||||
fixed_vfloat32m4_t exp_val =
|
||||
FP32Vec16(__riscv_vfneg_v_f32m4(
|
||||
__riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM))
|
||||
.exp()
|
||||
.reg;
|
||||
fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
|
||||
__riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
|
||||
return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
|
||||
}
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Type Traits & Global Helpers
|
||||
// ============================================================================
|
||||
|
||||
template <typename T>
|
||||
struct VecType {
|
||||
using vec_type = void;
|
||||
using vec_t = void;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using vec_t = typename VecType<T>::vec_type;
|
||||
|
||||
template <>
|
||||
struct VecType<float> {
|
||||
using vec_type = FP32Vec8;
|
||||
using vec_t = FP32Vec8;
|
||||
};
|
||||
template <>
|
||||
struct VecType<c10::Half> {
|
||||
using vec_type = FP16Vec8;
|
||||
using vec_t = FP16Vec8;
|
||||
};
|
||||
template <>
|
||||
struct VecType<c10::BFloat16> {
|
||||
using vec_type = BF16Vec8;
|
||||
using vec_t = BF16Vec8;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void storeFP32(float v, T* ptr) {
|
||||
*ptr = v;
|
||||
}
|
||||
template <>
|
||||
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
|
||||
*reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
|
||||
}
|
||||
|
||||
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
|
||||
reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
|
||||
}
|
||||
inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
|
||||
reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
|
||||
}
|
||||
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
|
||||
reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
|
||||
}
|
||||
inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
|
||||
acc = acc.fma(a, b);
|
||||
}
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
*ptr = static_cast<__bf16>(v);
|
||||
};
|
||||
inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
|
||||
: reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
|
||||
: reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
|
||||
#else
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
uint32_t val;
|
||||
std::memcpy(&val, &v, 4);
|
||||
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
|
||||
}
|
||||
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
|
||||
#endif
|
||||
|
||||
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
|
||||
|
||||
} // namespace vec_op
|
||||
|
||||
#ifndef CPU_KERNEL_GUARD_IN
|
||||
#define CPU_KERNEL_GUARD_IN(NAME)
|
||||
#endif
|
||||
|
||||
#ifndef CPU_KERNEL_GUARD_OUT
|
||||
#define CPU_KERNEL_GUARD_OUT(NAME)
|
||||
#endif
|
||||
|
||||
#endif // CPU_TYPES_RISCV_HPP
|
||||
#endif // CPU_TYPES_RISCV_HPP
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
#ifndef CPU_TYPES_RISCV_DEFS_HPP
|
||||
#define CPU_TYPES_RISCV_DEFS_HPP
|
||||
|
||||
// VLEN-to-LMUL mapping for RISC-V Vector extension.
|
||||
//
|
||||
// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
|
||||
// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
|
||||
// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
|
||||
|
||||
#include <riscv_vector.h>
|
||||
|
||||
#if __riscv_v_min_vlen == 128
|
||||
#define LMUL_128 m1
|
||||
#define LMUL_256 m2
|
||||
#define LMUL_512 m4
|
||||
#define LMUL_1024 m8
|
||||
#define BOOL_256 b16
|
||||
#define BOOL_512 b8
|
||||
#elif __riscv_v_min_vlen == 256
|
||||
#define LMUL_128 mf2
|
||||
#define LMUL_256 m1
|
||||
#define LMUL_512 m2
|
||||
#define LMUL_1024 m4
|
||||
#define BOOL_256 b32
|
||||
#define BOOL_512 b16
|
||||
#else
|
||||
#error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen"
|
||||
#endif
|
||||
|
||||
// Token-paste helpers.
|
||||
#define _RVV_P2(a, b) a##b
|
||||
#define _RVV_P3(a, b, c) a##b##c
|
||||
#define _RVV_P4(a, b, c, d) a##b##c##d
|
||||
#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
|
||||
#define RVVI(base, lmul) _RVV_P2(base, lmul)
|
||||
#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
|
||||
#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d)
|
||||
// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16
|
||||
#define _RVV_PB(base, lmul, btype) base##lmul##_##btype
|
||||
#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype)
|
||||
|
||||
// ---- Semantic fixed-vector typedefs (named by element count) ----
|
||||
|
||||
// float16
|
||||
typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
|
||||
// float32
|
||||
typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(1024)));
|
||||
|
||||
// int32
|
||||
typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
// uint16
|
||||
typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
// bfloat16
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
#endif
|
||||
|
||||
// ---- Reduction accumulator type (always m1 = one register of f32) ----
|
||||
// Used for scalar reductions; only element [0] is meaningful.
|
||||
typedef vfloat32m1_t rvv_f32_accum_t
|
||||
__attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen)));
|
||||
|
||||
// ---- Mask types for f32 elements ----
|
||||
#if __riscv_v_min_vlen == 128
|
||||
typedef vbool16_t rvv_mask_f32x8_t;
|
||||
typedef vbool8_t rvv_mask_f32x16_t;
|
||||
#elif __riscv_v_min_vlen == 256
|
||||
typedef vbool32_t rvv_mask_f32x8_t;
|
||||
typedef vbool16_t rvv_mask_f32x16_t;
|
||||
#endif
|
||||
|
||||
#endif // CPU_TYPES_RISCV_DEFS_HPP
|
||||
@@ -0,0 +1,905 @@
|
||||
#ifndef CPU_TYPES_RISCV_IMPL_HPP
|
||||
#define CPU_TYPES_RISCV_IMPL_HPP
|
||||
|
||||
// Shared implementation of RVV vector-type wrapper classes.
|
||||
// This file is VLEN-independent: it uses the semantic type names and
|
||||
// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp.
|
||||
//
|
||||
// DO NOT include this file directly; include cpu_types_riscv.hpp instead.
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <limits>
|
||||
#include <torch/all.h>
|
||||
namespace vec_op {
|
||||
|
||||
// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
|
||||
// is defined, otherwise via the FP32-simulation fallback path.
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
|
||||
|
||||
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
||||
|
||||
#define FORCE_INLINE __attribute__((always_inline)) inline
|
||||
|
||||
namespace {
|
||||
template <typename T, T... indexes, typename F>
|
||||
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
|
||||
(f(std::integral_constant<T, indexes>{}), ...);
|
||||
};
|
||||
} // namespace
|
||||
|
||||
template <typename T, T count, typename F,
|
||||
typename = std::enable_if_t<std::is_invocable_v<F, T>>>
|
||||
constexpr void unroll_loop(F&& f) {
|
||||
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct Vec {
|
||||
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
|
||||
};
|
||||
|
||||
struct FP32Vec8;
|
||||
struct FP32Vec16;
|
||||
|
||||
// ============================================================================
|
||||
// FP16 Implementation
|
||||
// ============================================================================
|
||||
|
||||
struct FP16Vec8 : public Vec<FP16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_fp16x8_t reg;
|
||||
|
||||
explicit FP16Vec8(const void* ptr)
|
||||
: reg(RVVI(__riscv_vle16_v_f16, LMUL_128)(
|
||||
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
|
||||
|
||||
explicit FP16Vec8(const FP32Vec8&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
|
||||
elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(_Float16);
|
||||
RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr),
|
||||
byte_stride, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_fp16x16_t reg;
|
||||
|
||||
explicit FP16Vec16(const void* ptr)
|
||||
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
|
||||
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16& vec);
|
||||
|
||||
void save(void* ptr) const {
|
||||
RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
|
||||
elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(_Float16);
|
||||
RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr),
|
||||
byte_stride, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// BF16 Implementation
|
||||
// ============================================================================
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
|
||||
FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
|
||||
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
|
||||
}
|
||||
FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) {
|
||||
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v);
|
||||
}
|
||||
FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) {
|
||||
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v);
|
||||
}
|
||||
|
||||
struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_bf16x8_t reg;
|
||||
|
||||
explicit BF16Vec8(const void* ptr)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16,
|
||||
LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {};
|
||||
explicit BF16Vec8(const FP32Vec8&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
|
||||
bf16_to_u16(reg), VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
|
||||
bf16_to_u16(reg), elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
|
||||
byte_stride, bf16_to_u16(reg),
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_bf16x16_t reg;
|
||||
|
||||
explicit BF16Vec16(const void* ptr)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16,
|
||||
LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
|
||||
bf16_to_u16(reg), VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
|
||||
bf16_to_u16(reg), elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
|
||||
byte_stride, bf16_to_u16(reg),
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
constexpr static int VEC_ELEM_NUM = 32;
|
||||
fixed_bf16x32_t reg;
|
||||
|
||||
explicit BF16Vec32(const void* ptr)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
|
||||
LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
|
||||
fixed_u16x32_t u16_combined =
|
||||
RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
|
||||
u16_val, u16_val, u16_val, u16_val);
|
||||
reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
|
||||
LMUL_512)(u16_combined);
|
||||
};
|
||||
|
||||
void save(void* ptr) const {
|
||||
RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
|
||||
bf16_to_u16(reg), VEC_ELEM_NUM);
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
|
||||
bf16_to_u16(reg), elem_num);
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
|
||||
byte_stride, bf16_to_u16(reg),
|
||||
VEC_ELEM_NUM);
|
||||
}
|
||||
};
|
||||
|
||||
#else
|
||||
// ============================================================================
|
||||
// BF16 Fallback Implementation (FP32 Simulation)
|
||||
// ============================================================================
|
||||
|
||||
struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_fp32x8_t reg_fp32;
|
||||
explicit BF16Vec8(const void* ptr) {
|
||||
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
|
||||
float tmp[8];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
|
||||
std::memcpy(&tmp[i], &v, 4);
|
||||
}
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
|
||||
}
|
||||
explicit BF16Vec8(const FP32Vec8&);
|
||||
void save(void* ptr) const {
|
||||
float tmp[8];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
float tmp[8];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < elem_num; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
float tmp[8];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
|
||||
uint8_t* u8 = static_cast<uint8_t*>(ptr);
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
uint16_t val = static_cast<uint16_t>(v >> 16);
|
||||
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_fp32x16_t reg_fp32;
|
||||
explicit BF16Vec16(const void* ptr) {
|
||||
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
|
||||
float tmp[16];
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
|
||||
std::memcpy(&tmp[i], &v, 4);
|
||||
}
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
|
||||
}
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
void save(void* ptr) const {
|
||||
float tmp[16];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save(void* ptr, int elem_num) const {
|
||||
float tmp[16];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < elem_num; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
float tmp[16];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
|
||||
uint8_t* u8 = static_cast<uint8_t*>(ptr);
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
uint16_t val = static_cast<uint16_t>(v >> 16);
|
||||
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
constexpr static int VEC_ELEM_NUM = 32;
|
||||
fixed_fp32x32_t reg_fp32;
|
||||
|
||||
explicit BF16Vec32(const void* ptr) {
|
||||
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
|
||||
float tmp[32];
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
|
||||
std::memcpy(&tmp[i], &v, 4);
|
||||
}
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
|
||||
}
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
float tmp_small[8];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
|
||||
float tmp_large[32];
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
|
||||
}
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32);
|
||||
}
|
||||
|
||||
void save(void* ptr) const {
|
||||
float tmp[32];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
|
||||
void save(void* ptr, int elem_num) const {
|
||||
float tmp[32];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
|
||||
uint16_t* u16 = static_cast<uint16_t*>(ptr);
|
||||
for (int i = 0; i < elem_num; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
u16[i] = static_cast<uint16_t>(v >> 16);
|
||||
}
|
||||
}
|
||||
|
||||
void save_strided(void* ptr, ptrdiff_t stride) const {
|
||||
float tmp[32];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
|
||||
uint8_t* u8 = static_cast<uint8_t*>(ptr);
|
||||
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
uint32_t v;
|
||||
std::memcpy(&v, &tmp[i], 4);
|
||||
uint16_t val = static_cast<uint16_t>(v >> 16);
|
||||
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
|
||||
}
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
// ============================================================================
|
||||
// FP32 Implementation
|
||||
// ============================================================================
|
||||
|
||||
struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
constexpr static int VEC_ELEM_NUM = 4;
|
||||
fixed_fp32x4_t reg;
|
||||
explicit FP32Vec4(float v)
|
||||
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec4()
|
||||
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec4(const float* ptr)
|
||||
: reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {};
|
||||
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
|
||||
void save(float* ptr) const {
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
void save(float* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num);
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
fixed_fp32x8_t reg;
|
||||
|
||||
explicit FP32Vec8(float v)
|
||||
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8()
|
||||
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(const float* ptr)
|
||||
: reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {};
|
||||
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
|
||||
explicit FP32Vec8(const FP16Vec8& v)
|
||||
: reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(fixed_fp16x8_t v)
|
||||
: reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
explicit FP32Vec8(fixed_bf16x8_t v)
|
||||
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(const BF16Vec8& v)
|
||||
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
|
||||
};
|
||||
#else
|
||||
explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
|
||||
#endif
|
||||
|
||||
float reduce_sum() const {
|
||||
rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
|
||||
scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar,
|
||||
VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
FP32Vec8 operator*(const FP32Vec8& b) const {
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 operator+(const FP32Vec8& b) const {
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 operator-(const FP32Vec8& b) const {
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 operator/(const FP32Vec8& b) const {
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 min(const FP32Vec8& b) const {
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 max(const FP32Vec8& b) const {
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec8 abs() const {
|
||||
return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
|
||||
return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num));
|
||||
}
|
||||
FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
|
||||
return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num));
|
||||
}
|
||||
|
||||
FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
|
||||
fixed_fp32x8_t temp =
|
||||
RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM);
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
void save(float* ptr) const {
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
void save(float* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num);
|
||||
}
|
||||
void save_strided(float* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(float);
|
||||
RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
|
||||
FP32Vec8 exp() const {
|
||||
// Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
|
||||
// Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
|
||||
// Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
|
||||
constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
|
||||
constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
|
||||
fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
const float inv_ln2 = 1.44269504088896341f;
|
||||
fixed_fp32x8_t x_scaled =
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM);
|
||||
fixed_i32x8_t n_int =
|
||||
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
|
||||
fixed_fp32x8_t n_float =
|
||||
RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x8_t r =
|
||||
RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x8_t poly =
|
||||
RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_i32x8_t biased_exp =
|
||||
RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM);
|
||||
biased_exp =
|
||||
RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM);
|
||||
fixed_i32x8_t exponent_bits =
|
||||
RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM);
|
||||
fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32,
|
||||
LMUL_256)(exponent_bits);
|
||||
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 tanh() const {
|
||||
fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
|
||||
VEC_ELEM_NUM);
|
||||
fixed_fp32x8_t x2 =
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM);
|
||||
FP32Vec8 exp_val = FP32Vec8(x2).exp();
|
||||
fixed_fp32x8_t num =
|
||||
RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
|
||||
fixed_fp32x8_t den =
|
||||
RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
|
||||
return FP32Vec8(
|
||||
RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec8 er() const {
|
||||
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
|
||||
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
|
||||
fixed_fp32x8_t abs_x =
|
||||
RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM);
|
||||
t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x8_t poly =
|
||||
RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(
|
||||
abs_x, abs_x, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM))
|
||||
.exp()
|
||||
.reg;
|
||||
fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)(
|
||||
reg, 0.0f, VEC_ELEM_NUM);
|
||||
return FP32Vec8(
|
||||
RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM));
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
fixed_fp32x16_t reg;
|
||||
|
||||
explicit FP32Vec16(float v)
|
||||
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16()
|
||||
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(const float* ptr)
|
||||
: reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {};
|
||||
explicit FP32Vec16(const FP32Vec8& data)
|
||||
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
|
||||
data.reg, data.reg)) {};
|
||||
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
explicit FP32Vec16(fixed_bf16x16_t v)
|
||||
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(const BF16Vec16& v)
|
||||
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) {
|
||||
};
|
||||
#else
|
||||
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
|
||||
#endif
|
||||
|
||||
FP32Vec16 operator+(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 operator-(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 operator*(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 operator/(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
float reduce_sum() const {
|
||||
rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
|
||||
scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar,
|
||||
VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
float reduce_max() const {
|
||||
rvv_f32_accum_t scalar =
|
||||
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
|
||||
scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar,
|
||||
VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
float reduce_min() const {
|
||||
rvv_f32_accum_t scalar =
|
||||
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
|
||||
scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar,
|
||||
VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
}
|
||||
|
||||
template <int group_size>
|
||||
float reduce_sub_sum(int idx) {
|
||||
static_assert(VEC_ELEM_NUM % group_size == 0);
|
||||
const int start = idx * group_size;
|
||||
auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM);
|
||||
rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)(
|
||||
RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start,
|
||||
VEC_ELEM_NUM),
|
||||
RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)(
|
||||
indices, start + group_size, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM);
|
||||
rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
|
||||
scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)(
|
||||
mask, reg, scalar, VEC_ELEM_NUM);
|
||||
return __riscv_vfmv_f_s_f32m1_f32(scalar);
|
||||
};
|
||||
|
||||
FP32Vec16 max(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 min(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
}
|
||||
FP32Vec16 abs() const {
|
||||
return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
|
||||
return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)(
|
||||
max_v.reg,
|
||||
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
void save(float* ptr) const {
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
void save(float* ptr, int elem_num) const {
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num);
|
||||
}
|
||||
void save_strided(float* ptr, ptrdiff_t stride) const {
|
||||
ptrdiff_t byte_stride = stride * sizeof(float);
|
||||
RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM);
|
||||
}
|
||||
|
||||
FP32Vec16 exp() const {
|
||||
// Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
|
||||
// Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
|
||||
// Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
|
||||
constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
|
||||
constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
|
||||
fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
const float inv_ln2 = 1.44269504088896341f;
|
||||
fixed_fp32x16_t x_scaled =
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM);
|
||||
fixed_i32x16_t n_int =
|
||||
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
|
||||
fixed_fp32x16_t n_float =
|
||||
RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM);
|
||||
fixed_fp32x16_t r =
|
||||
RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x16_t poly =
|
||||
RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
|
||||
0.009618129107628f, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
|
||||
0.055504108664821f, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
|
||||
0.240226506959101f, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
|
||||
0.693147180559945f, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)(
|
||||
RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0,
|
||||
VEC_ELEM_NUM);
|
||||
fixed_fp32x16_t scale =
|
||||
RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)(
|
||||
RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM));
|
||||
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 tanh() const {
|
||||
fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
|
||||
VEC_ELEM_NUM);
|
||||
FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)(
|
||||
x_clamped, 2.0f, VEC_ELEM_NUM))
|
||||
.exp();
|
||||
return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
|
||||
RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM));
|
||||
}
|
||||
|
||||
FP32Vec16 er() const {
|
||||
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
|
||||
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
|
||||
fixed_fp32x16_t abs_x =
|
||||
RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM);
|
||||
fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM),
|
||||
1.0f, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x16_t poly =
|
||||
RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1,
|
||||
VEC_ELEM_NUM);
|
||||
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM);
|
||||
|
||||
fixed_fp32x16_t exp_val =
|
||||
FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x,
|
||||
VEC_ELEM_NUM),
|
||||
VEC_ELEM_NUM))
|
||||
.exp()
|
||||
.reg;
|
||||
fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)(
|
||||
reg, 0.0f, VEC_ELEM_NUM);
|
||||
return FP32Vec16(
|
||||
RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM));
|
||||
}
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Type Traits & Global Helpers
|
||||
// ============================================================================
|
||||
|
||||
template <typename T>
|
||||
struct VecType {
|
||||
using vec_type = void;
|
||||
using vec_t = void;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using vec_t = typename VecType<T>::vec_type;
|
||||
|
||||
template <>
|
||||
struct VecType<float> {
|
||||
using vec_type = FP32Vec8;
|
||||
using vec_t = FP32Vec8;
|
||||
};
|
||||
template <>
|
||||
struct VecType<c10::Half> {
|
||||
using vec_type = FP16Vec8;
|
||||
using vec_t = FP16Vec8;
|
||||
};
|
||||
template <>
|
||||
struct VecType<c10::BFloat16> {
|
||||
using vec_type = BF16Vec8;
|
||||
using vec_t = BF16Vec8;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void storeFP32(float v, T* ptr) {
|
||||
*ptr = v;
|
||||
}
|
||||
template <>
|
||||
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
|
||||
*reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
|
||||
}
|
||||
|
||||
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
|
||||
reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM);
|
||||
}
|
||||
inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
|
||||
reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM);
|
||||
}
|
||||
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
|
||||
reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM);
|
||||
}
|
||||
inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
|
||||
acc = acc.fma(a, b);
|
||||
}
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
*ptr = static_cast<__bf16>(v);
|
||||
};
|
||||
inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
|
||||
: reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) {
|
||||
};
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
|
||||
: reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
|
||||
};
|
||||
#else
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
uint32_t val;
|
||||
std::memcpy(&val, &v, 4);
|
||||
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
|
||||
}
|
||||
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
|
||||
#endif
|
||||
|
||||
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
|
||||
|
||||
} // namespace vec_op
|
||||
|
||||
#ifndef CPU_KERNEL_GUARD_IN
|
||||
#define CPU_KERNEL_GUARD_IN(NAME)
|
||||
#endif
|
||||
|
||||
#ifndef CPU_KERNEL_GUARD_OUT
|
||||
#define CPU_KERNEL_GUARD_OUT(NAME)
|
||||
#endif
|
||||
|
||||
#endif // CPU_TYPES_RISCV_IMPL_HPP
|
||||
@@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size);
|
||||
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids);
|
||||
|
||||
namespace cpu_utils {
|
||||
void eagle_prepare_inputs_padded_kernel_impl(
|
||||
const torch::Tensor& cu_num_draft_tokens,
|
||||
@@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"block_size) -> ()",
|
||||
&compute_slot_mapping_kernel_impl);
|
||||
|
||||
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
|
||||
|
||||
// Speculative decoding kernels
|
||||
ops.def(
|
||||
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
|
||||
|
||||
+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) {
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
@@ -277,7 +277,9 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
(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) {
|
||||
// The shared-memory vectorized offset load only handles full 4-expert
|
||||
// chunks. Use the scalar specialization for the remainder cases.
|
||||
if (n_experts >= 4 && n_experts % 4 == 0) {
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
@@ -299,7 +301,9 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
n_experts);
|
||||
}
|
||||
} else {
|
||||
if (n_experts >= 16) {
|
||||
// The low-latency vectorized expert lookup only handles full 16-expert
|
||||
// chunks. Fall back to the scalar lookup path for the remainder cases.
|
||||
if (n_experts >= 16 && n_experts % 16 == 0) {
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
+2
-11
@@ -40,15 +40,6 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
|
||||
#define __HIP__FP8MFMA__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \
|
||||
defined(__gfx1150__) || defined(__gfx1151__))
|
||||
#define __HIP__GFX11__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX12__
|
||||
#endif
|
||||
|
||||
#if defined(NDEBUG)
|
||||
#undef NDEBUG
|
||||
#include <assert.h>
|
||||
@@ -1629,7 +1620,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
#elif defined(__HIP__GFX11__)
|
||||
#elif defined(__GFX11__)
|
||||
|
||||
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
|
||||
|
||||
@@ -2388,7 +2379,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
|
||||
out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
|
||||
}
|
||||
|
||||
#elif defined(__HIP__GFX12__)
|
||||
#elif defined(__GFX12__)
|
||||
|
||||
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
|
||||
|
||||
|
||||
+18
-23
@@ -26,16 +26,11 @@
|
||||
#define __HIP__GFX9__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && \
|
||||
(defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
|
||||
defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
|
||||
// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
|
||||
#if defined(__GFX11__) || defined(__GFX12__)
|
||||
#define __HIP__GFX1X__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX12__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
|
||||
#define __HIP__MI3XX__
|
||||
#endif
|
||||
@@ -1845,7 +1840,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
|
||||
return out_c;
|
||||
}
|
||||
|
||||
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
@@ -1893,7 +1888,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
|
||||
float sum[N][YTILE] = {};
|
||||
#else
|
||||
@@ -1931,7 +1926,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#pragma unroll
|
||||
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
#pragma unroll
|
||||
@@ -1955,7 +1950,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
}
|
||||
|
||||
// Final reduction
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
@@ -1993,7 +1988,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#endif
|
||||
|
||||
const bool writeback_lane =
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
threadIdx.x == (THRDS - 1);
|
||||
#else
|
||||
threadIdx.x == 0;
|
||||
@@ -2009,7 +2004,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
if (y + m >= M) break; // To avoid mem access fault.
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
float result = sum[n][y] * sA * sB;
|
||||
#else
|
||||
float result = sum[n][y][0] * sA * sB;
|
||||
@@ -2027,7 +2022,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
m += CuCount * _WvPrGrp * YTILE;
|
||||
}
|
||||
}
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
|
||||
@@ -2039,9 +2034,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
|
||||
const int _WvPrGrp, const int CuCount) {
|
||||
UNREACHABLE_CODE
|
||||
}
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
|
||||
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
@@ -2088,7 +2083,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
|
||||
float sum[N][YTILE] = {};
|
||||
#else
|
||||
@@ -2128,7 +2123,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#pragma unroll
|
||||
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
#pragma unroll
|
||||
@@ -2152,7 +2147,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
}
|
||||
|
||||
// Final reduction
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
@@ -2190,7 +2185,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#endif
|
||||
|
||||
const bool writeback_lane =
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
threadIdx.x == (THRDS - 1);
|
||||
#else
|
||||
threadIdx.x == 0;
|
||||
@@ -2206,7 +2201,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
if (y + m >= M) break; // To avoid mem access fault.
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
float result = sum[n][y] * sA * sB;
|
||||
#else
|
||||
float result = sum[n][y][0] * sA * sB;
|
||||
@@ -2224,7 +2219,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
m += CuCount * _WvPrGrp * YTILE;
|
||||
}
|
||||
}
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
|
||||
@@ -2236,7 +2231,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
|
||||
const int CuCount) {
|
||||
UNREACHABLE_CODE
|
||||
}
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
|
||||
void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
|
||||
const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
|
||||
|
||||
+6
-27
@@ -22,7 +22,7 @@
|
||||
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
|
||||
# =============================================================================
|
||||
|
||||
ARG CUDA_VERSION=13.0.0
|
||||
ARG CUDA_VERSION=13.0.2
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
@@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# Explicitly set the list to avoid issues with torch 2.2
|
||||
# See https://github.com/pytorch/pytorch/pull/123243
|
||||
# From versions.json: .torch.cuda_arch_list
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
#################### BUILD BASE IMAGE ####################
|
||||
|
||||
@@ -580,33 +580,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||
# https://docs.flashinfer.ai/installation.html
|
||||
# From versions.json: .flashinfer.version
|
||||
# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
|
||||
# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
|
||||
# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
|
||||
ARG FLASHINFER_VERSION=0.6.7
|
||||
ARG FLASHINFER_VERSION=0.6.8.post1
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
|
||||
&& flashinfer show-config
|
||||
|
||||
# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
|
||||
# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
|
||||
# which fails without internet. This step caches them at build time.
|
||||
RUN python3 <<'PYEOF'
|
||||
from flashinfer.jit import env as jit_env
|
||||
from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
|
||||
from flashinfer.artifacts import ArtifactPath, CheckSumHash
|
||||
|
||||
download_trtllm_headers(
|
||||
'bmm',
|
||||
jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
|
||||
f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
|
||||
ArtifactPath.TRTLLM_GEN_BMM,
|
||||
get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
|
||||
)
|
||||
|
||||
print('FlashInfer TRTLLM BMM headers downloaded successfully')
|
||||
PYEOF
|
||||
&& flashinfer show-config \
|
||||
&& flashinfer download-cubin
|
||||
|
||||
# ============================================================
|
||||
# OPENAI API SERVER DEPENDENCIES
|
||||
@@ -786,7 +765,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# install kv_connectors if requested
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system numba==0.61.2
|
||||
uv pip install --system numba==0.65.0
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
@@ -217,16 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
|
||||
|
||||
|
||||
# build flashinfer for torch nightly from source around 10 mins
|
||||
# release version: v0.6.7
|
||||
# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
|
||||
# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
|
||||
# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
|
||||
# release version: v0.6.8.post1
|
||||
# todo(elainewy): cache flashinfer build result for faster build
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
echo "git clone flashinfer..." \
|
||||
&& git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \
|
||||
&& git clone --depth 1 --branch v0.6.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
|
||||
&& cd flashinfer \
|
||||
&& git submodule update --init --recursive \
|
||||
&& echo "finish git clone flashinfer..." \
|
||||
|
||||
+74
-16
@@ -2,6 +2,11 @@
|
||||
ARG REMOTE_VLLM="0"
|
||||
ARG COMMON_WORKDIR=/app
|
||||
ARG BASE_IMAGE=rocm/vllm-dev:base
|
||||
# AMD NIC backend
|
||||
ARG NIC_BACKEND=none
|
||||
# AMD AINIC apt repo settings
|
||||
ARG AINIC_VERSION=1.117.5
|
||||
ARG UBUNTU_CODENAME=jammy
|
||||
|
||||
# Sccache configuration (only used in release pipeline)
|
||||
ARG USE_SCCACHE
|
||||
@@ -119,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="f33a5599"
|
||||
ARG RIXL_BRANCH="bf4a7214"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="da3fac2a"
|
||||
ARG UCX_REPO="https://github.com/ROCm/ucx.git"
|
||||
ARG UCX_BRANCH="7009d7a1"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
@@ -160,7 +165,7 @@ RUN cd /usr/local/src && \
|
||||
--disable-doxygen-doc \
|
||||
--enable-optimizations \
|
||||
--enable-devel-headers \
|
||||
--with-rocm=/opt/rocm \
|
||||
--with-rocm=${ROCM_PATH} \
|
||||
--with-verbs \
|
||||
--with-dm \
|
||||
--enable-mt && \
|
||||
@@ -181,7 +186,12 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
|
||||
ninja install
|
||||
|
||||
# Generate RIXL wheel
|
||||
RUN cd /opt/rixl && mkdir -p /app/install && \
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
@@ -192,9 +202,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 +213,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.
|
||||
@@ -217,6 +226,47 @@ RUN git clone ${DEEPEP_REPO} \
|
||||
&& git checkout ${DEEPEP_BRANCH} \
|
||||
&& python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
|
||||
|
||||
# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
|
||||
# not force users to rebuild the long-lived Dockerfile.rocm_base image.
|
||||
FROM base AS mori_base
|
||||
ARG NIC_BACKEND
|
||||
ARG AINIC_VERSION
|
||||
ARG UBUNTU_CODENAME
|
||||
RUN /bin/bash -lc 'set -euo pipefail; \
|
||||
echo "[MORI] Install MoRI proxy deps"; \
|
||||
pip install --quiet --ignore-installed blinker && \
|
||||
pip install --quiet quart msgpack aiohttp pyzmq; \
|
||||
echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
|
||||
\
|
||||
# NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
|
||||
# Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
|
||||
case "${NIC_BACKEND}" in \
|
||||
# default: mlx5
|
||||
none) \
|
||||
;; \
|
||||
# AMD NIC
|
||||
ainic) \
|
||||
apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
|
||||
rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
|
||||
curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
|
||||
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
|
||||
> /etc/apt/sources.list.d/amdainic.list; \
|
||||
apt-get update && apt-get install -y --no-install-recommends \
|
||||
libionic-dev \
|
||||
ionic-common \
|
||||
; \
|
||||
rm -rf /var/lib/apt/lists/*; \
|
||||
;; \
|
||||
# TODO: Add Broadcom bnxt packages/repos here later.
|
||||
# bnxt) \
|
||||
# echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
|
||||
# ;; \
|
||||
*) \
|
||||
echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
|
||||
exit 2; \
|
||||
;; \
|
||||
esac;'
|
||||
|
||||
# -----------------------
|
||||
# vLLM wheel release build stage (for building distributable wheels)
|
||||
# This stage pins dependencies to custom ROCm wheel versions and handles version detection
|
||||
@@ -319,7 +369,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# -----------------------
|
||||
# Test vLLM image
|
||||
FROM base AS test
|
||||
FROM mori_base AS test
|
||||
|
||||
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
@@ -386,6 +436,10 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
|
||||
ENV MIOPEN_DEBUG_CONV_DIRECT=0
|
||||
ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
|
||||
# Source code is used in the `python_only_compile.sh` test
|
||||
# We hide it inside `src/` so that this source code
|
||||
# will not be imported by other tests
|
||||
@@ -410,7 +464,7 @@ RUN printf '%s\n' \
|
||||
|
||||
# -----------------------
|
||||
# Final vLLM image
|
||||
FROM base AS final
|
||||
FROM mori_base AS final
|
||||
|
||||
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
@@ -447,6 +501,8 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
ARG BASE_IMAGE
|
||||
ARG NIC_BACKEND
|
||||
ARG AINIC_VERSION
|
||||
|
||||
# Copy over the benchmark scripts as well
|
||||
COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
|
||||
@@ -464,7 +520,9 @@ ENV HIP_FORCE_DEV_KERNARG=1
|
||||
# Workaround for ROCm profiler limits
|
||||
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
|
||||
ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
|
||||
RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt
|
||||
RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
|
||||
&& echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
|
||||
&& echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.10.post3"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="2d02c6a9"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
# Sccache configuration (only used in release pipeline)
|
||||
|
||||
@@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
@@ -164,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.6.0
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -20,7 +20,7 @@ variable "NVCC_THREADS" {
|
||||
}
|
||||
|
||||
variable "TORCH_CUDA_ARCH_LIST" {
|
||||
default = "8.0 8.9 9.0 10.0"
|
||||
default = "8.0 8.9 9.0 10.0 11.0 12.0"
|
||||
}
|
||||
|
||||
variable "COMMIT" {
|
||||
@@ -88,7 +88,6 @@ target "test-ubuntu2404" {
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
@@ -100,7 +99,6 @@ target "openai-ubuntu2404" {
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
|
||||
"variable": {
|
||||
"CUDA_VERSION": {
|
||||
"default": "13.0.0"
|
||||
"default": "13.0.2"
|
||||
},
|
||||
"PYTHON_VERSION": {
|
||||
"default": "3.12"
|
||||
@@ -11,10 +11,10 @@
|
||||
"default": "22.04"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
|
||||
},
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
|
||||
},
|
||||
"GET_PIP_URL": {
|
||||
"default": "https://bootstrap.pypa.io/get-pip.py"
|
||||
@@ -32,7 +32,7 @@
|
||||
"default": "false"
|
||||
},
|
||||
"TORCH_CUDA_ARCH_LIST": {
|
||||
"default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
|
||||
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX"
|
||||
},
|
||||
"MAX_JOBS": {
|
||||
"default": "2"
|
||||
@@ -65,7 +65,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.7"
|
||||
"default": "0.6.8.post1"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 315 KiB After Width: | Height: | Size: 315 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 108 KiB |
File diff suppressed because one or more lines are too long
@@ -108,6 +108,38 @@ P99 ITL (ms): 8.39
|
||||
==================================================
|
||||
```
|
||||
|
||||
#### Results Visualization
|
||||
|
||||
The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis.
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--endpoint /v1/completions \
|
||||
--dataset-name sharegpt \
|
||||
--dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
|
||||
--num-prompts 100 \
|
||||
--plot-timeline \
|
||||
--timeline-itl-thresholds 2,5 \
|
||||
--plot-dataset-stats \
|
||||
--save-result
|
||||
```
|
||||
|
||||
##### Interactive Timeline
|
||||
|
||||
The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms)
|
||||
|
||||
Example output:
|
||||
|
||||
<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
|
||||
|
||||
##### Dataset statistics
|
||||
|
||||
The generated figure shows the input prompt and output tokens distribution.
|
||||
|
||||
Example output: 
|
||||
|
||||
#### Custom Dataset
|
||||
|
||||
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
|
||||
|
||||
@@ -155,9 +155,9 @@ switch to `--physcpubind=<cpu-list> --membind=<node>`.
|
||||
|
||||
These `--numa-bind*` options only apply to GPU execution processes. They do not
|
||||
configure the CPU backend's separate thread-affinity controls. Automatic
|
||||
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based platforms;
|
||||
other GPU backends must provide explicit binding lists if they use these
|
||||
options.
|
||||
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based as well as
|
||||
ROCM-based platforms; other GPU backends must provide explicit binding lists if
|
||||
they use these options.
|
||||
|
||||
`--numa-bind-nodes` takes one non-negative NUMA node index per visible GPU, in
|
||||
the same order as the GPU indices.
|
||||
|
||||
@@ -66,7 +66,7 @@ This is for controlling general behavior of the API when serving your model:
|
||||
|
||||
See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls.
|
||||
|
||||
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server passes you the resampled waveform and task parameters; you return a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
|
||||
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
|
||||
|
||||
#### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n)
|
||||
|
||||
@@ -75,21 +75,20 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt
|
||||
??? code "get_generation_prompt()"
|
||||
|
||||
```python
|
||||
from vllm.config.speech_to_text import SpeechToTextParams
|
||||
|
||||
class YourASRModel(nn.Module, SupportsTranscription):
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(
|
||||
cls,
|
||||
audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: str | None,
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: str | None,
|
||||
stt_params: SpeechToTextParams,
|
||||
) -> PromptType:
|
||||
# Example with a free-form instruction prompt
|
||||
audio = stt_params.audio
|
||||
stt_config = stt_params.stt_config
|
||||
task_type = stt_params.task_type
|
||||
|
||||
task_word = "Transcribe" if task_type == "transcribe" else "Translate"
|
||||
prompt = (
|
||||
"<start_of_turn>user\n"
|
||||
@@ -112,20 +111,22 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries:
|
||||
??? code "get_generation_prompt()"
|
||||
|
||||
```python
|
||||
from vllm.config.speech_to_text import SpeechToTextParams
|
||||
|
||||
class YourASRModel(nn.Module, SupportsTranscription):
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(
|
||||
cls,
|
||||
audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: str | None,
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: str | None,
|
||||
stt_params: SpeechToTextParams,
|
||||
) -> PromptType:
|
||||
audio = stt_params.audio
|
||||
stt_config = stt_params.stt_config
|
||||
language = stt_params.language
|
||||
task_type = stt_params.task_type
|
||||
request_prompt = stt_params.request_prompt
|
||||
|
||||
if language is None:
|
||||
raise ValueError("Language must be specified")
|
||||
|
||||
@@ -193,7 +194,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,22 +207,20 @@ 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))
|
||||
prompts = []
|
||||
for chunk in chunks:
|
||||
prompt = self.model_cls.get_generation_prompt(
|
||||
stt_params = request.build_stt_params(
|
||||
audio=chunk,
|
||||
stt_config=self.asr_config,
|
||||
model_config=self.model_config,
|
||||
language=language,
|
||||
task_type=self.task_type,
|
||||
request_prompt=request.prompt,
|
||||
to_language=to_language,
|
||||
)
|
||||
prompt = self.model_cls.get_generation_prompt(stt_params)
|
||||
prompts.append(prompt)
|
||||
return prompts, duration
|
||||
```
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
|
||||
|
||||
- [Deployment with CPUs](#deployment-with-cpus)
|
||||
- [Deployment with GPUs](#deployment-with-gpus)
|
||||
- [Serving with gRPC](#serving-with-grpc)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"](#startup-probe-or-readiness-probe-failure-container-log-contains-keyboardinterrupt-terminated)
|
||||
- [Conclusion](#conclusion)
|
||||
@@ -387,6 +388,49 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
|
||||
|
||||
If the service is correctly deployed, you should receive a response from the vLLM model.
|
||||
|
||||
## Serving with gRPC
|
||||
|
||||
vLLM can serve models over gRPC instead of HTTP by passing the `--grpc` flag. This requires the optional gRPC dependencies:
|
||||
|
||||
```bash
|
||||
pip install vllm[grpc]
|
||||
```
|
||||
|
||||
When using `--grpc`, the server exposes the standard [gRPC Health Checking Protocol](https://github.com/grpc/grpc/blob/master/doc/health-checking.md) (`grpc.health.v1.Health`), which integrates with Kubernetes [native gRPC probes](https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#define-a-grpc-liveness-probe) (available since Kubernetes 1.24).
|
||||
|
||||
To deploy with gRPC, change the `vllm serve` command to include `--grpc` and replace `httpGet` probes with `grpc` probes:
|
||||
|
||||
```yaml
|
||||
containers:
|
||||
- name: mistral-7b
|
||||
image: vllm/vllm-openai:latest
|
||||
command: ["/bin/sh", "-c"]
|
||||
args: [
|
||||
"pip install vllm[grpc] && vllm serve mistralai/Mistral-7B-Instruct-v0.3 --grpc --port 50051 --trust-remote-code"
|
||||
]
|
||||
ports:
|
||||
- containerPort: 50051
|
||||
livenessProbe:
|
||||
grpc:
|
||||
port: 50051
|
||||
initialDelaySeconds: 120
|
||||
periodSeconds: 10
|
||||
readinessProbe:
|
||||
grpc:
|
||||
port: 50051
|
||||
initialDelaySeconds: 120
|
||||
periodSeconds: 5
|
||||
```
|
||||
|
||||
!!! note
|
||||
The gRPC health service checks the engine status on every probe. If the engine is unhealthy or the server is shutting down, the probe returns `NOT_SERVING`.
|
||||
|
||||
You can also verify the health service manually with `grpcurl`:
|
||||
|
||||
```bash
|
||||
grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -167,10 +169,10 @@ Priority is **1 = highest** (tried first).
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
@@ -213,7 +215,7 @@ configuration.
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
|
||||
@@ -52,14 +52,14 @@ For each graph replay:
|
||||
|
||||
When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
|
||||
|
||||
### Video inference support (experimental)
|
||||
### Video inference support
|
||||
|
||||
Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
|
||||
|
||||
!!! note
|
||||
Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
|
||||
|
||||
Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
|
||||
Mixed inputs (image+video) per prompt are also supported now.
|
||||
|
||||
## Model integration via `SupportsEncoderCudaGraph`
|
||||
|
||||
@@ -76,6 +76,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
* `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
|
||||
* `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
|
||||
* `get_input_modality(...)` - return the modality of the inputs.
|
||||
* `get_max_frames_per_video()` - return model-specific max frames per video.
|
||||
|
||||
!!! note
|
||||
The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
|
||||
@@ -96,7 +97,7 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
|
||||
* `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
|
||||
* `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
|
||||
* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
|
||||
* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
|
||||
* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value according to its `processing_info`). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
|
||||
|
||||
## Usage guide
|
||||
|
||||
@@ -142,7 +143,6 @@ Enable encoder CUDA Graphs via `compilation_config`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-VL-32B \
|
||||
--limit-mm-per-prompt '{"image": 0}' \
|
||||
--compilation-config '{"cudagraph_mm_encoder": true}'
|
||||
```
|
||||
|
||||
@@ -150,7 +150,6 @@ With explicit budgets:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-VL-32B \
|
||||
--limit-mm-per-prompt '{"image": 0}' \
|
||||
--compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
|
||||
```
|
||||
|
||||
@@ -169,7 +168,6 @@ compilation_config = {
|
||||
|
||||
model = vllm.LLM(
|
||||
model="Qwen/Qwen3-VL-32B",
|
||||
limit_mm_per_prompt='{"image": 0}',
|
||||
compilation_config=compilation_config,
|
||||
)
|
||||
```
|
||||
|
||||
+76
-2
@@ -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 |
|
||||
@@ -30,6 +31,7 @@ or just on the low or high end.
|
||||
| [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
|
||||
| [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
|
||||
| [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always |
|
||||
| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm → single kernel | O1 (ROCm/AITER only) | ~2% | No | Always |
|
||||
|
||||
## Support Matrix
|
||||
|
||||
@@ -40,8 +42,9 @@ 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_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
|
||||
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
|
||||
| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — |
|
||||
| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
|
||||
@@ -49,11 +52,15 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
| `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group |
|
||||
| `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group |
|
||||
| `fuse_act_padding` | — | — | — | — | FP16/BF16 |
|
||||
| `fuse_mla_dual_rms_norm` | — | — | — | — | BF16 |
|
||||
|
||||
\* `fuse_attn_quant` support depends on the attention backend in use; not all backends support
|
||||
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
|
||||
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`.
|
||||
@@ -145,7 +152,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt
|
||||
|
||||
- `FLASHINFER`: CUDA sm100+ with FlashInfer installed
|
||||
|
||||
`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
|
||||
`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`
|
||||
|
||||
The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
|
||||
with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
|
||||
@@ -184,6 +191,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,
|
||||
@@ -347,6 +383,44 @@ when the hidden size is 2880 and AITER Triton GEMMs *not* enabled.
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`RocmAiterTritonAddRMSNormPadFusionPass`)
|
||||
|
||||
### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`)
|
||||
|
||||
!!! info
|
||||
ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention.
|
||||
|
||||
!!! note
|
||||
When the native implementation of `rms_norm` is used (the default on CUDA and
|
||||
ROCm for now), Inductor's built-in fusion already handles merging these norms
|
||||
automatically. This explicit pass targets the case where AITER's custom
|
||||
`rms_norm` op is active, which Inductor cannot fuse on its own.
|
||||
|
||||
**What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm
|
||||
operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER,
|
||||
reducing kernel launch overhead from 2 launches to 1 per MLA layer.
|
||||
|
||||
```text
|
||||
# Unfused:
|
||||
q_c, kv_lora = split(projected, [q_dim, kv_dim])
|
||||
kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
|
||||
q_c = rms_norm(q_c, q_weight, eps)
|
||||
kv_c = rms_norm(kv_c, kv_weight, eps)
|
||||
|
||||
# Fused:
|
||||
q_c, kv_lora = split(projected, [q_dim, kv_dim])
|
||||
kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
|
||||
q_normed, kv_normed = fused_mla_dual_rms_norm(
|
||||
q_c, q_weight, kv_c, kv_weight, eps1, eps2)
|
||||
```
|
||||
|
||||
Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above
|
||||
when AITER is available.
|
||||
|
||||
**Code locations.**
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)
|
||||
- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)
|
||||
- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)
|
||||
|
||||
## See Also
|
||||
|
||||
- [Optimization Levels](optimization_levels.md) — high-level presets that set
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -83,8 +83,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
|
||||
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
|
||||
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
|
||||
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
|
||||
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
|
||||
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
|
||||
|
||||
@@ -56,6 +56,7 @@ Fusions:
|
||||
- `-cc.pass_config.fuse_norm_quant=True`*
|
||||
- `-cc.pass_config.fuse_act_quant=True`*
|
||||
- `-cc.pass_config.fuse_act_padding=True`†
|
||||
- `-cc.pass_config.fuse_mla_dual_rms_norm=True`†
|
||||
|
||||
\* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.</br>
|
||||
† These fusions are ROCm-only and require AITER.
|
||||
|
||||
@@ -104,7 +104,7 @@ for output in outputs:
|
||||
Batch invariance has been tested and verified on the following models:
|
||||
|
||||
- **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
|
||||
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`
|
||||
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
|
||||
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
# Context Extension
|
||||
|
||||
!!! note
|
||||
The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead.
|
||||
This directory contains examples for extending the context length of models using vLLM.
|
||||
|
||||
## Offline Inference Example
|
||||
|
||||
The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
|
||||
|
||||
### Usage
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/context_extension.py
|
||||
```
|
||||
|
||||
## OpenAI Online Method
|
||||
|
||||
You can also use vLLM's OpenAI-compatible API to serve models with extended context length.
|
||||
|
||||
### Usage
|
||||
|
||||
Run the vLLM server with the following command to extend the context length using YARN:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \
|
||||
--max-model-len 131072
|
||||
```
|
||||
|
||||
### Client Example
|
||||
|
||||
After starting the server, you can use the OpenAI Python client to interact with it:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:8000/v1",
|
||||
api_key="token-abc123" # Dummy API key, required by the client
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
max_tokens=128,
|
||||
temperature=0.8,
|
||||
top_p=0.95
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
### Key Parameters
|
||||
|
||||
The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters).
|
||||
|
||||
Common parameters include:
|
||||
|
||||
- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic")
|
||||
- `factor`: The factor by which to extend the context length
|
||||
- `original_max_position_embeddings`: The original maximum position embeddings of the model
|
||||
|
||||
The following parameters are specific to vLLM:
|
||||
|
||||
- `max_model_len`: The new maximum sequence length after extension (original * factor).
|
||||
Used for KV cache pre‑allocation and request limit at serving time.
|
||||
@@ -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(
|
||||
@@ -780,6 +780,70 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
|
||||
Works with common video formats like MP4 when using OpenCV backends.
|
||||
|
||||
#### Pre-extracted Frame Sequences with `media_io_kwargs`
|
||||
|
||||
When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction.
|
||||
|
||||
**Supported Parameters:**
|
||||
|
||||
| Parameter | Type | Description |
|
||||
| --------- | ---- | ----------- |
|
||||
| `fps` | float | Frame rate of the original video |
|
||||
| `frames_indices` | list[int] | Indices of the actually sampled frames |
|
||||
| `total_num_frames` | int | Total frame count of the original video |
|
||||
| `duration` | float | Duration of the original video in seconds |
|
||||
| `do_sample_frames` | bool | Whether to perform frame sampling |
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
|
||||
|
||||
# Client-side frame extraction
|
||||
frames = extract_frames(video_path, num_frames=32)
|
||||
frames_b64 = ",".join([encode_image(f) for f in frames])
|
||||
video_url = f"data:video/jpeg;base64,{frames_b64}"
|
||||
|
||||
# Pass video metadata via media_io_kwargs
|
||||
response = client.chat.completions.create(
|
||||
model="your-multimodal-model",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "video_url", "video_url": {"url": video_url}},
|
||||
{"type": "text", "text": "Describe what happens in this video."}
|
||||
]
|
||||
}],
|
||||
extra_body={
|
||||
"media_io_kwargs": {
|
||||
"video": {
|
||||
"fps": 30.0,
|
||||
"frames_indices": [0, 10, 20, 30, 40, 50, 60, 70, 80, 90,
|
||||
100, 110, 120, 130, 140, 150, 160, 170,
|
||||
180, 190, 200, 210, 220, 230, 240, 250,
|
||||
260, 270, 280, 290, 300, 310],
|
||||
"total_num_frames": 900,
|
||||
"duration": 30.0,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
**Why use `media_io_kwargs`?**
|
||||
|
||||
When extracting frames client-side, the server loses important context about the original video:
|
||||
|
||||
- **Temporal information**: Which frames were sampled and their positions in the original timeline
|
||||
- **Video duration**: How long the original video was
|
||||
- **Frame rate**: The original playback speed
|
||||
|
||||
By passing this metadata, the model can better understand the temporal distribution of the sampled frames and whether important moments might have been skipped.
|
||||
|
||||
#### Custom RGBA Background Color
|
||||
|
||||
To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
|
||||
@@ -832,7 +896,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)
|
||||
|
||||
|
||||
@@ -17,6 +17,9 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
|
||||
|
||||
# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
|
||||
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
|
||||
|
||||
# MXFP8 quantization for weights and activations
|
||||
llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8")
|
||||
```
|
||||
|
||||
Or with the CLI:
|
||||
@@ -24,6 +27,7 @@ Or with the CLI:
|
||||
```bash
|
||||
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
|
||||
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
|
||||
vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
|
||||
```
|
||||
|
||||
## Supported Schemes
|
||||
@@ -32,8 +36,7 @@ vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
|
||||
| ------ | ------------- | ------------------ | ----- |
|
||||
| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
|
||||
| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
|
||||
|
||||
Support for additional schemes will be added in future versions of vllm.
|
||||
| `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback |
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
|
||||
@@ -283,9 +283,7 @@ curl http://localhost:8000/v1/chat/completions \
|
||||
"messages": [
|
||||
{ "role": "user", "content": "9.11 and 9.8, which is greater?" }
|
||||
],
|
||||
"extra_body": {
|
||||
"thinking_token_budget": 10
|
||||
}
|
||||
"thinking_token_budget": 10
|
||||
}'
|
||||
```
|
||||
|
||||
|
||||
@@ -35,6 +35,99 @@ For reproducible measurements in your environment, use
|
||||
[`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py)
|
||||
or the [benchmark CLI guide](../../benchmarking/cli.md).
|
||||
|
||||
## `--speculative-config` schema
|
||||
|
||||
Use `--speculative-config` to pass speculative decoding settings as a JSON
|
||||
object on the CLI:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "draft_model",
|
||||
"model": "<draft-model>",
|
||||
"num_speculative_tokens": 5
|
||||
}'
|
||||
```
|
||||
|
||||
The same keys are accepted from Python via `LLM(..., speculative_config={...})`.
|
||||
The tables below highlight common user-facing keys accepted in this JSON
|
||||
object; they are not an exhaustive schema reference.
|
||||
For more details, see the generated [engine arguments reference](../../configuration/engine_args.md)
|
||||
and the API docs for [vllm.config.SpeculativeConfig][].
|
||||
|
||||
### Common keys
|
||||
|
||||
These keys are commonly used across speculative decoding setups, though some
|
||||
only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
|
||||
`dflash`.
|
||||
|
||||
| Key | Type | Default | Allowed values / meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. |
|
||||
| `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. |
|
||||
| `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. |
|
||||
| `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. |
|
||||
| `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. |
|
||||
| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
|
||||
| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
|
||||
| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
|
||||
|
||||
### Method-specific keys
|
||||
|
||||
#### N-gram
|
||||
|
||||
| Key | Type | Default | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `prompt_lookup_max` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_min` when omitted | Maximum n-gram window size. |
|
||||
| `prompt_lookup_min` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_max` when omitted | Minimum n-gram window size. |
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 4,
|
||||
"prompt_lookup_min": 2,
|
||||
"prompt_lookup_max": 5
|
||||
}'
|
||||
```
|
||||
|
||||
#### Suffix decoding
|
||||
|
||||
| Key | Type | Default | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `suffix_decoding_max_tree_depth` | `integer` | `24` | Maximum combined prefix-match and speculation tree depth. |
|
||||
| `suffix_decoding_max_cached_requests` | `integer` | `10000` | Maximum number of requests cached in the global suffix tree. Set `0` to disable the global cache. |
|
||||
| `suffix_decoding_max_spec_factor` | `float` | `1.0` | Caps speculative length as a multiple of prefix-match length. |
|
||||
| `suffix_decoding_min_token_prob` | `float` | `0.1` | Minimum estimated token probability required to speculate a token. |
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 8,
|
||||
"suffix_decoding_max_tree_depth": 24,
|
||||
"suffix_decoding_max_cached_requests": 10000,
|
||||
"suffix_decoding_max_spec_factor": 1.0,
|
||||
"suffix_decoding_min_token_prob": 0.1
|
||||
}'
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--speculative-config` expects a JSON object on the CLI. In YAML config
|
||||
files, use a nested mapping instead of an escaped JSON string.
|
||||
- `tensor_parallel_size` is not a valid key in `speculative_config`. Use
|
||||
`draft_tensor_parallel_size` instead.
|
||||
- Keys such as `temperature` and `top_p` are sampling parameters, not
|
||||
`--speculative-config` fields.
|
||||
- Internal fields such as `target_model_config`, `draft_model_config`,
|
||||
`target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
|
||||
are populated by vLLM and are not intended to be set by users.
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
|
||||
@@ -33,9 +33,9 @@ vllm serve Qwen/Qwen3-4B-Thinking-2507 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
```
|
||||
|
||||
The code used to request as completions as a client remains unchanged:
|
||||
@@ -77,4 +77,8 @@ The code used to request as completions as a client remains unchanged:
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
|
||||
Note: Please use `--speculative-config` to set all configurations related
|
||||
to speculative decoding. The previous method of specifying the model
|
||||
through `--speculative-model` and adding related parameters such as
|
||||
`--num-speculative-tokens` separately has been deprecated. For supported
|
||||
keys and examples, see the [`--speculative-config` schema](README.md#--speculative-config-schema).
|
||||
|
||||
@@ -38,7 +38,7 @@ for output in outputs:
|
||||
```bash
|
||||
vllm serve XiaomiMiMo/MiMo-7B-Base \
|
||||
--tensor-parallel-size 1 \
|
||||
--speculative_config '{"method":"mtp","num_speculative_tokens":1}'
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
@@ -36,9 +36,9 @@ vllm serve Qwen/Qwen3-4B \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--speculative-config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
|
||||
```
|
||||
|
||||
## Pre-trained PARD weights
|
||||
|
||||
@@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
|
||||
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg CUDA_VERSION=13.0.1 \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--build-arg max_jobs=256 \
|
||||
--build-arg nvcc_threads=2 \
|
||||
--build-arg RUN_WHEEL_CHECK=false \
|
||||
|
||||
@@ -46,7 +46,7 @@ pip install -v -r requirements/xpu.txt
|
||||
|
||||
!!! note
|
||||
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
|
||||
- For torch 2.10 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.6.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
- For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
|
||||
- Finally, build and install vLLM XPU backend:
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr
|
||||
|-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
|
||||
| `classify` (see note) | Sequence-wise | reranker score for each sequence | `cross-encoder` | linear classifier |
|
||||
| `embed` | Sequence-wise | vector representations for each sequence | `bi-encoder` | cosine similarity |
|
||||
| `token_classify` | Token-wise | probability vector of classes for each token | nan | nan |
|
||||
| `token_classify` | Token-wise | probability vector of classes for each token | N/A | N/A |
|
||||
| `token_embed` | Token-wise | vector representations for each token | `late-interaction` | late interaction(MaxSim) |
|
||||
|
||||
!!! note
|
||||
@@ -86,14 +86,15 @@ The scoring models is designed to compute similarity scores between two input pr
|
||||
|
||||
### Pooling Usages
|
||||
|
||||
| Pooling Usages | Description |
|
||||
|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
|
||||
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
|
||||
| Token Classification Usages | Token-wise classification |
|
||||
| Token Embedding Usages | Token-wise embedding |
|
||||
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
|
||||
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
|
||||
| Pooling Usages | Description |
|
||||
|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
|
||||
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
|
||||
| Token Classification Usages | Token-wise classification |
|
||||
| Token Embedding Usages | Token-wise embedding |
|
||||
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
|
||||
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
|
||||
| Plugins Usages | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |
|
||||
|
||||
We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs.
|
||||
|
||||
@@ -101,9 +102,9 @@ For more detailed information, please refer to the link below.
|
||||
|
||||
- [Classification Usages](classify.md)
|
||||
- [Embedding Usages](embed.md)
|
||||
- [Reward Usages](reward.md)
|
||||
- [Token Classification Usages](token_classify.md)
|
||||
- [Token Embedding Usages](token_embed.md)
|
||||
- [Reward Usages](reward.md)
|
||||
- [Scoring Usages](scoring.md)
|
||||
- [Specific Model Examples](specific_models.md)
|
||||
|
||||
@@ -113,15 +114,17 @@ Each pooling model in vLLM supports one or more of these tasks according to
|
||||
[Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
|
||||
enabling the corresponding APIs.
|
||||
|
||||
### Offline APIs corresponding to pooling tasks
|
||||
### Offline APIs corresponding to pooling usages
|
||||
|
||||
| Task | APIs |
|
||||
|------------------|---------------------------------------------------------------------------------------|
|
||||
| `embed` | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) |
|
||||
| `classify` | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)` |
|
||||
| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")` |
|
||||
| `token_embed` | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)` |
|
||||
| `plugin` | `LLM.encode(..., pooling_task="plugin")` |
|
||||
| Pooling Usages | Dedicated API | Pooling task for `LLM.encode` API | Score Types | scoring function |
|
||||
|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------|
|
||||
| Classification Usages | `LLM.classify(...)` | `classify` | `cross-encoder` (see note) | linear classifier |
|
||||
| Embedding Usages | `LLM.embed(...)` | `embed` | `bi-encoder` | cosine similarity |
|
||||
| Token Classification Usages | N/A | `token_classify` | N/A | N/A |
|
||||
| Token Embedding Usages | N/A | `token_embed` | `late-interaction` | late interaction(MaxSim) |
|
||||
| Reward Usages | N/A | `classify` & `token_classify` | N/A | N/A |
|
||||
| Scoring Usages | `LLM.score(...)` | N/A | N/A | N/A |
|
||||
| Plugins Usages | N/A | `plugin` | N/A | N/A |
|
||||
|
||||
!!! note
|
||||
Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
@@ -147,7 +150,7 @@ It is primarily designed for [score models](scoring.md).
|
||||
|
||||
The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
|
||||
|
||||
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
|
||||
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
|
||||
|
||||
### Examples
|
||||
|
||||
@@ -183,9 +186,12 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all
|
||||
|
||||
The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.
|
||||
|
||||
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
|
||||
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
|
||||
|
||||
Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py)
|
||||
Code examples:
|
||||
|
||||
- [Online example](../../../examples/pooling/reward/token_reward_online.py)
|
||||
- [Offline example](../../../examples/pooling/reward/token_reward_offline.py)
|
||||
|
||||
### Examples
|
||||
|
||||
@@ -286,10 +292,10 @@ Pooling models now support token-wise task.
|
||||
|
||||
### Score task
|
||||
|
||||
`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
|
||||
classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
|
||||
equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
|
||||
### Pooling multitask support
|
||||
|
||||
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
|
||||
Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
|
||||
you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
|
||||
|
||||
@@ -134,3 +134,13 @@ print(f"Data: {data!r}")
|
||||
## Online Serving
|
||||
|
||||
Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary).
|
||||
|
||||
## More examples
|
||||
|
||||
More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward)
|
||||
|
||||
## Deprecated Features
|
||||
|
||||
### `LLM.reward`
|
||||
|
||||
`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead.
|
||||
|
||||
@@ -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. | ✅︎ | ✅︎ |
|
||||
@@ -418,6 +419,7 @@ th {
|
||||
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
|
||||
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
|
||||
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
|
||||
| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
|
||||
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
|
||||
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -471,6 +473,7 @@ th {
|
||||
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
|
||||
| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
|
||||
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
|
||||
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
|
||||
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -554,10 +557,12 @@ 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. | ✅︎ | ✅︎ |
|
||||
| `GlmOcrForConditionalGeneration` | GLM-OCR | T + I<sup>E+</sup> | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `Granite4VisionForConditionalGeneration` | Granite 4 Vision | T + I<sup>E+</sup> | `ibm-granite/granite-4.1-3b-vision`, etc. | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
|
||||
| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
|
||||
| `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | |
|
||||
@@ -633,6 +638,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 +649,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 +685,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.
|
||||
|
||||
@@ -153,6 +153,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
|
||||
| `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
|
||||
| `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
|
||||
| `policy` | The policy type for expert parallel load balancing | `"default"` |
|
||||
| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` |
|
||||
|
||||
For example:
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
+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
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user