forked from Karylab-cklius/vllm
Compare commits
219
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
efa494e397 | ||
|
|
74b77a593a | ||
|
|
ec53889a3b | ||
|
|
ede3d4ddf6 | ||
|
|
7301a834fa | ||
|
|
7104ac6b5d | ||
|
|
9a98eccb49 | ||
|
|
2ce1e2ba71 | ||
|
|
2496f66f7f | ||
|
|
d18b7f2723 | ||
|
|
80b633f076 | ||
|
|
28f589ebb1 | ||
|
|
98998279d4 | ||
|
|
71d208089e | ||
|
|
0ac6fc8352 | ||
|
|
fdfb2566c0 | ||
|
|
8a5cf1ccd6 | ||
|
|
fe1d923afc | ||
|
|
32daf56b42 | ||
|
|
82a42234be | ||
|
|
af9f583344 | ||
|
|
bd2d83ff31 | ||
|
|
bb78168b21 | ||
|
|
89c6a41001 | ||
|
|
7fdfa6441d | ||
|
|
e9b728de8a | ||
|
|
5828a205ef | ||
|
|
7a74f31d2e | ||
|
|
47930b59ca | ||
|
|
6aec99f030 | ||
|
|
f4966f8b3d | ||
|
|
2c9c07c85e | ||
|
|
320c52b134 | ||
|
|
6deb05e0e4 | ||
|
|
d82ac00923 | ||
|
|
dac9e9a640 | ||
|
|
d7607ad273 | ||
|
|
d955745d58 | ||
|
|
e1ed89dbee | ||
|
|
1c2ffc6f88 | ||
|
|
ca4cfd8731 | ||
|
|
c9c1540e61 | ||
|
|
c1d754d681 | ||
|
|
01d8cd92dd | ||
|
|
a4b14b98c6 | ||
|
|
cf1c906724 | ||
|
|
766ce2bb6b | ||
|
|
3d119f78f7 | ||
|
|
1b1359c332 | ||
|
|
cad4ca12b8 | ||
|
|
b697119800 | ||
|
|
b4c6dc6454 | ||
|
|
2ee5106372 | ||
|
|
7a89b72564 | ||
|
|
dc10e467a9 | ||
|
|
ee4d7df2b5 | ||
|
|
3e8afdf785 | ||
|
|
6690a0c4de | ||
|
|
1c23c42030 | ||
|
|
b12e42d132 | ||
|
|
69fdaffbcd | ||
|
|
80e2c4462d | ||
|
|
5b3807e862 | ||
|
|
59401ac9f1 | ||
|
|
d841386d27 | ||
|
|
fff9210b2a | ||
|
|
70db1488c5 | ||
|
|
2385e140d6 | ||
|
|
dab60fc658 | ||
|
|
996222f4bf | ||
|
|
e6fc848d4f | ||
|
|
f843ac1a1c | ||
|
|
7c2aa3108a | ||
|
|
ebf53ba373 | ||
|
|
baacbfcebf | ||
|
|
d8218b1ee7 | ||
|
|
9f153aa781 | ||
|
|
d3de61502f | ||
|
|
540aaf2140 | ||
|
|
4128605ad4 | ||
|
|
e2f993dc41 | ||
|
|
05cb606cad | ||
|
|
3f627ebef7 | ||
|
|
bc941f375d | ||
|
|
6afa25000c | ||
|
|
823a0ab754 | ||
|
|
2c27c294c0 | ||
|
|
ba94a3b998 | ||
|
|
dc68bd8c41 | ||
|
|
753e9d55e6 | ||
|
|
ac3409d162 | ||
|
|
93ee4cd47f | ||
|
|
980796cd07 | ||
|
|
5add018beb | ||
|
|
d5fe994e79 | ||
|
|
fa662b1a8b | ||
|
|
3c0b4432be | ||
|
|
469f3dcf1d | ||
|
|
94fcdd007f | ||
|
|
d9ff7e4e9a | ||
|
|
967c5c3bc3 | ||
|
|
54c660c3a6 | ||
|
|
8fb0274415 | ||
|
|
eebce65756 | ||
|
|
303916e93d | ||
|
|
5633405964 | ||
|
|
6124a98a9b | ||
|
|
2ed0a9627b | ||
|
|
4dcd10eb0d | ||
|
|
228bcc436b | ||
|
|
3d3ba460a2 | ||
|
|
66ecfd0568 | ||
|
|
f0f6805d8a | ||
|
|
15652a6b70 | ||
|
|
51ef688831 | ||
|
|
6ac69203e8 | ||
|
|
1505b3d8a1 | ||
|
|
32f34d3935 | ||
|
|
9c7f7741d4 | ||
|
|
6181e80fe0 | ||
|
|
3bb46975bd | ||
|
|
810966453a | ||
|
|
2a983c79ac | ||
|
|
bc5745a00f | ||
|
|
3b3d5287fa | ||
|
|
062b05ff3a | ||
|
|
fa27d4e9cf | ||
|
|
67d3792d99 | ||
|
|
00d1fb7747 | ||
|
|
c9b4b184b4 | ||
|
|
f87df1df9e | ||
|
|
eafbb06331 | ||
|
|
ec0a31d4aa | ||
|
|
c8beda4cc3 | ||
|
|
2f27c9a150 | ||
|
|
4765f0f189 | ||
|
|
a50e675b0d | ||
|
|
f6a708ab2b | ||
|
|
4200f62147 | ||
|
|
c73b0d0db9 | ||
|
|
e28e369f78 | ||
|
|
703fb17b13 | ||
|
|
b593396c7a | ||
|
|
91e17d4315 | ||
|
|
aa6fb8a329 | ||
|
|
6a894574bf | ||
|
|
7f003a1285 | ||
|
|
ef0df7dbd6 | ||
|
|
a80af24356 | ||
|
|
c66b19800b | ||
|
|
6a11d72df7 | ||
|
|
02d2da0748 | ||
|
|
bbb6c274c8 | ||
|
|
62215e72c6 | ||
|
|
7fe7800fa4 | ||
|
|
8a83e6f2d7 | ||
|
|
efc347f1b2 | ||
|
|
d98b8f371c | ||
|
|
e64237ae82 | ||
|
|
d61d8566ec | ||
|
|
d2f70da116 | ||
|
|
6542d48964 | ||
|
|
ca73293fa6 | ||
|
|
ef3af56a97 | ||
|
|
b4a6f26c90 | ||
|
|
165b7864d0 | ||
|
|
c505cd93ef | ||
|
|
96229fa99e | ||
|
|
da1daf40bf | ||
|
|
4efd6ffde0 | ||
|
|
56aff0dd15 | ||
|
|
063ce98fb7 | ||
|
|
62d6f06e3d | ||
|
|
b7c5baf63d | ||
|
|
a55fccfc7c | ||
|
|
41a4829f22 | ||
|
|
38fd2405f3 | ||
|
|
a947f7a420 | ||
|
|
439203d32c | ||
|
|
8d9536a775 | ||
|
|
3da29aa4a5 | ||
|
|
06f94633e7 | ||
|
|
99ef652907 | ||
|
|
4cc78c9d5d | ||
|
|
3dbb4e0ace | ||
|
|
b21443e23c | ||
|
|
06ee2d8433 | ||
|
|
b5235fca2e | ||
|
|
3e77036768 | ||
|
|
6f68ca3e91 | ||
|
|
0c96dd64fb | ||
|
|
68f5e565c9 | ||
|
|
9354fb1ba5 | ||
|
|
f35b557239 | ||
|
|
e68988a248 | ||
|
|
4b87b3e845 | ||
|
|
90619351e3 | ||
|
|
d0975a4b50 | ||
|
|
1bdc60ed53 | ||
|
|
a6183563b6 | ||
|
|
22c2e87555 | ||
|
|
d01d0b4646 | ||
|
|
b4b4aaa70e | ||
|
|
5e2af28838 | ||
|
|
4f423bd5bc | ||
|
|
f0cd590d62 | ||
|
|
e6018c644a | ||
|
|
f25952e59b | ||
|
|
b58e082d95 | ||
|
|
0c1e6f63f5 | ||
|
|
ceb0111a90 | ||
|
|
0414d75410 | ||
|
|
128adabfe0 | ||
|
|
bdbf08fc02 | ||
|
|
6bad553f4e | ||
|
|
91945b6e4a | ||
|
|
2b237c7a41 | ||
|
|
dad95e34d8 | ||
|
|
a248b45d05 |
@@ -28,18 +28,19 @@ steps:
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
|
||||
- label: CPU-Compatibility Tests
|
||||
depends_on: []
|
||||
device: intel_cpu
|
||||
no_plugin: true
|
||||
source_file_dependencies:
|
||||
- cmake/cpu_extension.cmake
|
||||
- setup.py
|
||||
- vllm/platforms/cpu.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
|
||||
# Note: SDE can't be downloaded from CI host because of AWS WAF
|
||||
# - label: CPU-Compatibility Tests
|
||||
# depends_on: []
|
||||
# device: intel_cpu
|
||||
# no_plugin: true
|
||||
# source_file_dependencies:
|
||||
# - cmake/cpu_extension.cmake
|
||||
# - setup.py
|
||||
# - vllm/platforms/cpu.py
|
||||
# commands:
|
||||
# - |
|
||||
# bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
|
||||
# bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
|
||||
|
||||
- label: CPU-Language Generation and Pooling Model Tests
|
||||
depends_on: []
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
group: Basic Correctness
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: XPU Sleep Mode
|
||||
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/basic_correctness/test_cumem.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
pytest -v -s basic_correctness/test_mem.py::test_end_to_end'
|
||||
@@ -0,0 +1,23 @@
|
||||
group: Expert Parallelism
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: EPLB Algorithm
|
||||
key: eplb-algorithm
|
||||
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/distributed/eplb
|
||||
- tests/distributed/test_eplb_algo.py
|
||||
- tests/distributed/test_eplb_utils.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s distributed/test_eplb_algo.py'
|
||||
@@ -38,7 +38,17 @@ steps:
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/logger.py
|
||||
- vllm/model_executor/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/v1/sample
|
||||
- tests/v1/logits_processors
|
||||
- tests/v1/test_oracle.py
|
||||
@@ -52,4 +62,126 @@ steps:
|
||||
pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
|
||||
pytest -v -s v1/test_oracle.py &&
|
||||
pytest -v -s v1/test_request.py &&
|
||||
pytest -v -s v1/test_outputs.py'
|
||||
pytest -v -s v1/test_outputs.py &&
|
||||
pytest -v -s v1/sample/test_topk_topp_sampler.py'
|
||||
|
||||
- label: XPU CPU Offload
|
||||
timeout_in_minutes: 60
|
||||
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/
|
||||
- vllm/v1/kv_offload/
|
||||
- vllm/v1/kv_connector/
|
||||
- tests/v1/kv_offload/
|
||||
- tests/v1/kv_connector/unit/test_offloading_connector.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/kv_offload &&
|
||||
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
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/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/test_regression
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install modelscope &&
|
||||
cd tests &&
|
||||
pytest -v -s test_regression.py'
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs)
|
||||
key: metrics-tracing-2-gpus
|
||||
timeout_in_minutes: 30
|
||||
num_devices: 2
|
||||
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/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/tracing/
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/v1/tracing
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install opentelemetry-sdk\>=1.26.0 opentelemetry-api\>=1.26.0 opentelemetry-exporter-otlp\>=1.26.0 opentelemetry-semantic-conventions-ai\>=0.4.1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/tracing'
|
||||
|
||||
- label: Async Engine, Inputs, Utils, Worker
|
||||
key: async-engine-inputs-utils-worker
|
||||
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/assets/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/tokenizers/
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pip install av &&
|
||||
pytest -v -s detokenizer &&
|
||||
pytest -v -s -m "not cpu_test" ./multimodal &&
|
||||
pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
group: Models - Multimodal
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
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/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
|
||||
pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
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/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
|
||||
key: multi-modal-models-standard-3-llava-qwen2-vl
|
||||
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/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
|
||||
pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
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/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
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/
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
pip install open-clip-torch --no-deps &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
--deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]"
|
||||
--deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]"
|
||||
--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
|
||||
parallelism: 4
|
||||
@@ -40,7 +40,9 @@ steps:
|
||||
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
|
||||
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
|
||||
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
|
||||
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
|
||||
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
|
||||
VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
|
||||
python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
|
||||
'
|
||||
- label: "XPU V1 test"
|
||||
depends_on:
|
||||
@@ -83,5 +85,5 @@ steps:
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av &&
|
||||
cd tests &&
|
||||
pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
|
||||
pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
|
||||
pytest -v -s benchmarks/test_serve_cli.py'
|
||||
|
||||
@@ -28,8 +28,10 @@
|
||||
###############################################################################
|
||||
set -o pipefail
|
||||
|
||||
# Export Python path
|
||||
export PYTHONPATH=".."
|
||||
# Export Python path for commands that run directly on the host. Containerized
|
||||
# tests set this to /vllm-workspace below so spawned Python processes do not
|
||||
# depend on their current working directory.
|
||||
export PYTHONPATH="${PYTHONPATH:-..}"
|
||||
|
||||
###############################################################################
|
||||
# Helper Functions
|
||||
@@ -377,6 +379,14 @@ HF_CACHE="$(realpath ~)/huggingface"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
# Hugging Face Hub defaults to 10s request/download timeouts, while the ROCm
|
||||
# CI image currently raises downloads to 60s. AMD model-test jobs routinely
|
||||
# start from a cold or partially-populated shared cache, and the 60s read cap
|
||||
# has still timed out before pytest reached the vLLM behavior under test.
|
||||
# Keep the CI default explicit and overridable from the Buildkite environment.
|
||||
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
|
||||
: "${HF_HUB_ETAG_TIMEOUT:=60}"
|
||||
|
||||
# ---- Command source selection ----
|
||||
# Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
|
||||
# Fall back to $* for backward compatibility, but warn that inner
|
||||
@@ -416,7 +426,14 @@ fi
|
||||
|
||||
echo "Final commands: $commands"
|
||||
|
||||
MYPYTHONPATH=".."
|
||||
MYPYTHONPATH="/vllm-workspace"
|
||||
|
||||
container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}"
|
||||
container_job_id="${container_job_id//[^A-Za-z0-9_.-]/_}"
|
||||
container_job_id_short="${container_job_id:0:8}"
|
||||
CONTAINER_TMPDIR="/tmp/vllm-${container_job_id_short}"
|
||||
CONTAINER_CACHE_ROOT="/tmp/vllm-buildkite-${container_job_id}/cache"
|
||||
CONTAINER_PREFLIGHT="mkdir -p \"\$TMPDIR\" \"\$TORCHINDUCTOR_CACHE_DIR\" \"\$TRITON_CACHE_DIR\" \"\$VLLM_CACHE_ROOT\" \"\$XDG_CACHE_HOME\" && python -c \"import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing\""
|
||||
|
||||
# Verify GPU access
|
||||
render_gid=$(getent group render | cut -d: -f3)
|
||||
@@ -493,6 +510,8 @@ else
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
-e HF_TOKEN \
|
||||
-e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
|
||||
-e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-e BUILDKITE_PARALLEL_JOB \
|
||||
@@ -500,10 +519,15 @@ else
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
-e "TMPDIR=${CONTAINER_TMPDIR}/tmp" \
|
||||
-e "TORCHINDUCTOR_CACHE_DIR=${CONTAINER_CACHE_ROOT}/torchinductor" \
|
||||
-e "TRITON_CACHE_DIR=${CONTAINER_CACHE_ROOT}/triton" \
|
||||
-e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
|
||||
-e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
|
||||
-e "PYTORCH_ROCM_ARCH=" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${commands}"
|
||||
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
|
||||
|
||||
exit_code=$?
|
||||
handle_pytest_exit "$exit_code"
|
||||
|
||||
@@ -324,23 +324,6 @@ IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
|
||||
|
||||
echo "Using image: ${IMAGE}"
|
||||
|
||||
if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
|
||||
echo "Image already exists locally, skipping pull"
|
||||
else
|
||||
echo "Image not found locally, waiting for lock..."
|
||||
|
||||
flock /tmp/docker-pull.lock bash -c "
|
||||
if docker image inspect '${IMAGE}' >/dev/null 2>&1; then
|
||||
echo 'Image already pulled by another runner'
|
||||
else
|
||||
echo 'Pulling image...'
|
||||
timeout 900 docker pull '${IMAGE}'
|
||||
fi
|
||||
"
|
||||
|
||||
echo "Pull step completed"
|
||||
fi
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
}
|
||||
@@ -357,9 +340,12 @@ export HF_TOKEN ZE_AFFINITY_MASK
|
||||
|
||||
{
|
||||
flock 9
|
||||
if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
|
||||
echo 'Image missing before container creation, pulling again...'
|
||||
if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
|
||||
echo "Image already exists locally, skipping pull"
|
||||
else
|
||||
echo "Image not found locally, pulling image..."
|
||||
timeout 900 docker pull "${IMAGE}"
|
||||
echo "Pull step completed"
|
||||
fi
|
||||
|
||||
docker create \
|
||||
@@ -372,6 +358,8 @@ export HF_TOKEN ZE_AFFINITY_MASK
|
||||
--entrypoint='' \
|
||||
-e HF_TOKEN \
|
||||
-e ZE_AFFINITY_MASK \
|
||||
-e BUILDKITE_PARALLEL_JOB \
|
||||
-e BUILDKITE_PARALLEL_JOB_COUNT \
|
||||
-e CMDS \
|
||||
--name "${container_name}" \
|
||||
"${IMAGE}" \
|
||||
|
||||
@@ -110,6 +110,36 @@ install_uv() {
|
||||
| env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
|
||||
}
|
||||
|
||||
setup_pyo3_python() {
|
||||
local python_version="${PYO3_PYTHON_VERSION:-3.12}"
|
||||
|
||||
log_section "Installing Python ${python_version} for PyO3 tests"
|
||||
uv python install "$python_version"
|
||||
PYO3_PYTHON="$(uv python find \
|
||||
--managed-python \
|
||||
--no-project \
|
||||
--resolve-links \
|
||||
"$python_version")"
|
||||
export PYO3_PYTHON
|
||||
|
||||
local python_libdir
|
||||
python_libdir="$("$PYO3_PYTHON" - <<'PY'
|
||||
import pathlib
|
||||
import sysconfig
|
||||
|
||||
libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
|
||||
ldlibrary = sysconfig.get_config_var("LDLIBRARY")
|
||||
assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
|
||||
assert ldlibrary
|
||||
assert (libdir / ldlibrary).exists(), libdir / ldlibrary
|
||||
print(libdir)
|
||||
PY
|
||||
)"
|
||||
|
||||
export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
|
||||
export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
|
||||
}
|
||||
|
||||
run_style_clippy() {
|
||||
install_cargo_sort
|
||||
|
||||
@@ -132,6 +162,7 @@ run_style_clippy() {
|
||||
|
||||
run_tests() {
|
||||
install_uv
|
||||
setup_pyo3_python
|
||||
install_cargo_nextest
|
||||
|
||||
log_section "Running cargo nextest"
|
||||
|
||||
+207
-129
@@ -88,16 +88,16 @@
|
||||
# - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. #
|
||||
# * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. #
|
||||
# * [Plugin Tests (2 GPUs)]: #
|
||||
# - {`pytest -v -s entrypoints/openai/test_oot_registration.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s models/test_oot_registration.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s plugins/lora_resolvers`}: Unit tests for in-tree lora resolver plugins #
|
||||
# - {`pytest -v -s plugins_tests/test_oot_registration_online.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s plugins_tests/test_oot_registration_offline.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s plugins_tests/lora_resolvers`}: Unit tests for in-tree lora resolver plugins #
|
||||
# * [LoRA TP (Distributed)]: #
|
||||
# - There is some Tensor Parallelism related processing logic in LoRA that requires multi-GPU testing for validation. #
|
||||
# - {`pytest -v -s -x lora/test_gptoss_tp.py`}: Disabled for now because MXFP4 backend on non-cuda platform doesn't support #
|
||||
# LoRA yet. #
|
||||
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Don't test llama model here, it seems hf implementation is buggy. See: #
|
||||
# https://github.com/vllm-project/vllm/pull/5689 #
|
||||
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 #
|
||||
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 #
|
||||
# in favor of new tests in fusions_e2e. We avoid replicating the new jobs in #
|
||||
# this file as it's deprecated. #
|
||||
# #
|
||||
@@ -315,24 +315,6 @@ steps:
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py
|
||||
|
||||
#---------------------------------------------------------- mi250 · engine -----------------------------------------------------------#
|
||||
|
||||
- label: Engine # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/engine
|
||||
- tests/test_sequence
|
||||
- tests/test_config
|
||||
- tests/test_logger
|
||||
- tests/test_vllm_port
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
|
||||
|
||||
#----------------------------------------------------------- mi250 · evals -----------------------------------------------------------#
|
||||
|
||||
- label: Multi-Modal Accuracy Eval (Small Models) # TBD
|
||||
@@ -635,9 +617,9 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
- pytest -v -s models/test_oot_registration.py
|
||||
- pytest -v -s plugins/lora_resolvers
|
||||
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------#
|
||||
|
||||
@@ -822,7 +804,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: V1 e2e (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -848,7 +830,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------#
|
||||
|
||||
@@ -908,10 +890,10 @@ steps:
|
||||
- vllm/
|
||||
- tests/basic_correctness/test_basic_correctness
|
||||
- tests/basic_correctness/test_cpu_offload
|
||||
- tests/basic_correctness/test_cumem.py
|
||||
- tests/basic_correctness/test_mem.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s basic_correctness/test_cumem.py
|
||||
- pytest -v -s basic_correctness/test_mem.py
|
||||
- pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
@@ -1205,7 +1187,39 @@ steps:
|
||||
|
||||
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
|
||||
|
||||
- label: Entrypoints Integration (API Server 2) # TBD
|
||||
- label: Entrypoints Unit Tests # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/unit_tests
|
||||
- tests/entrypoints/weight_transfer
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s entrypoints/unit_tests
|
||||
- pytest -v -s entrypoints/weight_transfer
|
||||
|
||||
- label: Entrypoints Integration (LLM) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
|
||||
- label: Entrypoints Integration (API Server) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
@@ -1221,7 +1235,7 @@ steps:
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
@@ -1235,90 +1249,45 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
- pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
|
||||
- label: Entrypoints Integration (API Server Generate) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/tool_parsers
|
||||
- tests/entrypoints/anthropic
|
||||
- tests/entrypoints/generate
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (LLM) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
- tests/entrypoints/offline_mode
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/llm/test_generate.py
|
||||
- pytest -v -s entrypoints/offline_mode
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
- pytest -v -s entrypoints/tool_parsers
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
|
||||
- label: Entrypoints Integration (Responses API) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1334,29 +1303,57 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Entrypoints Unit Tests # TBD
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: OpenAI API correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
@@ -1410,6 +1407,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -2184,10 +2182,72 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: Speculators Correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/speculative.py
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/transformers_utils/config.py
|
||||
- vllm/transformers_utils/configs/speculators/
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/v1/worker/gpu_model_runner.py
|
||||
- vllm/v1/sample/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/model_executor/models/llama_eagle3.py
|
||||
- vllm/model_executor/models/qwen3.py
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- vllm/model_executor/models/registry.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/evals/gsm8k/
|
||||
- tests/v1/spec_decode/test_speculators_correctness.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
|
||||
- label: Extract Hidden States Integration # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/speculative.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/
|
||||
- vllm/model_executor/layers/attention/
|
||||
- vllm/model_executor/layers/mamba/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/model_executor/models/extract_hidden_states.py
|
||||
- vllm/model_executor/models/llama.py
|
||||
- vllm/model_executor/models/qwen3_5.py
|
||||
- vllm/model_executor/models/qwen3_next.py
|
||||
- vllm/model_executor/models/registry.py
|
||||
- vllm/transformers_utils/configs/extract_hidden_states.py
|
||||
- vllm/transformers_utils/configs/qwen3_5.py
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/v1/kv_cache_interface.py
|
||||
- vllm/v1/spec_decode/extract_hidden_states.py
|
||||
- vllm/v1/worker/gpu_model_runner.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: V1 attention (H100-MI300) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
@@ -2346,7 +2406,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Distributed DP Tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2382,7 +2442,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2396,7 +2456,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2410,7 +2470,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: V1 e2e (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2546,6 +2606,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2555,7 +2616,7 @@ steps:
|
||||
- tests/test_logger
|
||||
- tests/test_vllm_port
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
|
||||
|
||||
#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------#
|
||||
|
||||
@@ -2637,6 +2698,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2736,7 +2798,7 @@ steps:
|
||||
|
||||
#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
|
||||
|
||||
- label: Entrypoints Integration (API Server 2) # TBD
|
||||
- label: Entrypoints Integration (API Server) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2752,7 +2814,7 @@ steps:
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2766,9 +2828,9 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2780,29 +2842,31 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
- label: Entrypoints Integration (API Server Generate) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/tool_parsers
|
||||
- tests/entrypoints/anthropic
|
||||
- tests/entrypoints/generate
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/tool_parsers
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2818,6 +2882,20 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -2995,7 +3073,7 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
commands:
|
||||
- rocm-smi
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
|
||||
- pytest -v -s tests/kernels/attention/test_attention_selector.py
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
@@ -3359,7 +3437,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -3374,7 +3452,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -3389,7 +3467,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ group: Attention
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: V1 attention (H100)
|
||||
key: v1-attention-h100
|
||||
- label: V1 attention (H100-MI300)
|
||||
key: v1-attention-h100-mi300
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
@@ -13,6 +13,20 @@ steps:
|
||||
- tests/v1/attention
|
||||
commands:
|
||||
- pytest -v -s v1/attention
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 70
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
- vllm/model_executor/layers/attention
|
||||
- vllm/v1/attention
|
||||
- tests/v1/attention
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/envs.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: V1 attention (B200)
|
||||
key: v1-attention-b200
|
||||
|
||||
@@ -10,9 +10,9 @@ steps:
|
||||
- vllm/
|
||||
- tests/basic_correctness/test_basic_correctness
|
||||
- tests/basic_correctness/test_cpu_offload
|
||||
- tests/basic_correctness/test_cumem.py
|
||||
- tests/basic_correctness/test_mem.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s basic_correctness/test_cumem.py
|
||||
- pytest -v -s basic_correctness/test_mem.py
|
||||
- pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
@@ -26,6 +26,12 @@ steps:
|
||||
- tests/test_jit_monitor.py
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Engine (1 GPU)
|
||||
key: engine-1-gpu
|
||||
@@ -38,7 +44,7 @@ steps:
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -8,10 +8,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/
|
||||
- tests/entrypoints/unit_tests
|
||||
- tests/entrypoints/weight_transfer
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
- pytest -v -s entrypoints/unit_tests
|
||||
- pytest -v -s entrypoints/weight_transfer
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
@@ -20,84 +21,20 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
- tests/entrypoints/offline_mode
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1)
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2)
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3)
|
||||
key: entrypoints-integration-api-server-openai-part-3
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server 2)
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-2
|
||||
timeout_in_minutes: 130
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -109,11 +46,79 @@ steps:
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server Generate)
|
||||
key: entrypoints-integration-api-server-generate
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/tool_parsers
|
||||
- tests/entrypoints/anthropic
|
||||
- tests/entrypoints/generate
|
||||
commands:
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/tool_parsers
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-speech_to_text
|
||||
@@ -126,6 +131,18 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-multimodal
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
@@ -137,16 +154,6 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: OpenAI API Correctness
|
||||
key: openai-api-correctness
|
||||
timeout_in_minutes: 30
|
||||
@@ -156,3 +163,20 @@ steps:
|
||||
- vllm/entrypoints/openai/
|
||||
commands: # LMEval
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/model_executor/model_loader/
|
||||
commands:
|
||||
- bash ../tools/install_torchcodec_rocm.sh || exit 1
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
|
||||
@@ -21,8 +21,9 @@ steps:
|
||||
- csrc/
|
||||
- tests/kernels/core
|
||||
- tests/kernels/test_concat_mla_q.py
|
||||
- tests/kernels/test_fused_qk_norm_rope_gate.py
|
||||
commands:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py
|
||||
|
||||
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
|
||||
key: kernels-minimax-reduce-rms-test-2-gpus
|
||||
@@ -86,7 +87,7 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -299,3 +300,4 @@ steps:
|
||||
- vllm/config
|
||||
commands:
|
||||
- pytest -v -s kernels/moe/test_moe_layer.py
|
||||
- pytest -v -s kernels/moe/test_deepep_v2_moe.py
|
||||
|
||||
@@ -12,6 +12,21 @@ steps:
|
||||
autorun_on_main: true
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
- vllm/model_executor/models/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
# - label: LM Eval Large Models (4 GPUs)(A100)
|
||||
# device: a100
|
||||
|
||||
@@ -300,9 +300,9 @@ steps:
|
||||
- tests/multimodal
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
- tests/tokenizers_
|
||||
- tests/reasoning
|
||||
- tests/tool_parsers
|
||||
- tests/tokenizers_
|
||||
- tests/parser
|
||||
- tests/transformers_utils
|
||||
- tests/config
|
||||
@@ -315,9 +315,9 @@ steps:
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
|
||||
- pytest -v -s tool_parsers
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s parser
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
@@ -51,6 +51,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 90
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
|
||||
@@ -15,7 +15,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -33,7 +33,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -50,7 +50,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -153,3 +153,12 @@ steps:
|
||||
- tests/models/multimodal/pooling
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal/pooling
|
||||
|
||||
@@ -37,6 +37,6 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s models/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
@@ -37,6 +37,21 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/transformers_utils/configs/speculators/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Speculators + MTP Nightly B200
|
||||
key: spec-decode-speculators-mtp-nightly-b200
|
||||
@@ -61,6 +76,20 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Draft Model
|
||||
key: spec-decode-draft-model
|
||||
@@ -72,6 +101,20 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Draft Model Nightly B200
|
||||
key: spec-decode-draft-model-nightly-b200
|
||||
|
||||
@@ -33,10 +33,3 @@ share/python-wheels/
|
||||
*.egg
|
||||
MANIFEST
|
||||
rust/target/
|
||||
# Not needed in Docker builds
|
||||
docs/
|
||||
.github/
|
||||
.pre-commit-config.yaml
|
||||
.clang-format
|
||||
.gitattributes
|
||||
format.sh
|
||||
|
||||
+8
-2
@@ -23,9 +23,14 @@
|
||||
|
||||
# Any change to the VllmConfig changes can have a large user-facing impact,
|
||||
# so spam a lot of people
|
||||
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
|
||||
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
|
||||
/vllm/config/cache.py @heheda12345
|
||||
|
||||
# Config utils
|
||||
/vllm/config/utils.py @hmellor
|
||||
/vllm/engine/arg_utils.py @hmellor
|
||||
/vllm/utils/argparse_utils.py
|
||||
|
||||
# Entrypoints
|
||||
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
|
||||
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
|
||||
@@ -34,10 +39,11 @@
|
||||
/vllm/entrypoints/speech_to_text/realtime @njhill
|
||||
/vllm/entrypoints/speech_to_text @NickLucche
|
||||
/vllm/entrypoints/pooling @noooop
|
||||
/vllm/entrypoints/sagemaker @DarkLight1337
|
||||
/vllm/entrypoints/serve/sagemaker @DarkLight1337
|
||||
/vllm/entrypoints/serve @njhill
|
||||
/vllm/entrypoints/*.py @njhill
|
||||
/vllm/entrypoints/chat_utils.py @DarkLight1337
|
||||
/vllm/entrypoints/offline_utils.py @DarkLight1337
|
||||
/vllm/entrypoints/llm.py @DarkLight1337
|
||||
|
||||
# Rust Frontend
|
||||
|
||||
+3
-12
@@ -21,6 +21,9 @@ pull_request_rules:
|
||||
- check-failure=pre-commit
|
||||
- -closed
|
||||
- -draft
|
||||
- or:
|
||||
- label=ready
|
||||
- label=verified
|
||||
actions:
|
||||
comment:
|
||||
message: |
|
||||
@@ -36,18 +39,6 @@ pull_request_rules:
|
||||
|
||||
For future commits, `pre-commit` will run automatically on changed files before each commit.
|
||||
|
||||
> [!TIP]
|
||||
> <details>
|
||||
> <summary>Is <code>mypy</code> failing?</summary>
|
||||
> <br/>
|
||||
> <code>mypy</code> is run differently in CI. If the failure is related to this check, please use the following command to run it locally:
|
||||
>
|
||||
> ```bash
|
||||
> # For mypy (substitute "3.10" with the failing version if needed)
|
||||
> pre-commit run --hook-stage manual mypy-3.10
|
||||
> ```
|
||||
> </details>
|
||||
|
||||
- name: comment-dco-failure
|
||||
description: Comment on PR when DCO check fails
|
||||
conditions:
|
||||
|
||||
@@ -15,7 +15,7 @@ jobs:
|
||||
actions: write
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
|
||||
- uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0
|
||||
with:
|
||||
# Increasing this value ensures that changes to this workflow
|
||||
# propagate to all issues and PRs in days rather than months
|
||||
|
||||
+7
-13
@@ -148,33 +148,27 @@ repos:
|
||||
language: python
|
||||
entry: python tools/pre_commit/generate_nightly_torch_test.py
|
||||
files: ^requirements/test/cuda\.(in|txt)$
|
||||
- id: mypy-local
|
||||
name: Run mypy locally for lowest supported Python version
|
||||
entry: python tools/pre_commit/mypy.py 0 "3.10"
|
||||
stages: [pre-commit] # Don't run in CI
|
||||
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.10
|
||||
entry: python tools/pre_commit/mypy.py "3.10"
|
||||
<<: &mypy_common
|
||||
language: python
|
||||
types_or: [python, pyi]
|
||||
require_serial: true
|
||||
additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
|
||||
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.10
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.10"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
|
||||
- id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.11
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.11"
|
||||
entry: python tools/pre_commit/mypy.py "3.11"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
- id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.12
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.12"
|
||||
entry: python tools/pre_commit/mypy.py "3.12"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
- id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.13
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.13"
|
||||
entry: python tools/pre_commit/mypy.py "3.13"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
- id: shellcheck
|
||||
|
||||
@@ -98,11 +98,13 @@ pre-commit run --all-files
|
||||
pre-commit run ruff-check --all-files
|
||||
|
||||
# Run mypy as it is in CI:
|
||||
pre-commit run mypy-3.10 --all-files --hook-stage manual
|
||||
pre-commit run mypy-3.12 --all-files --hook-stage manual
|
||||
```
|
||||
|
||||
The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
|
||||
|
||||
Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
|
||||
|
||||
### Commit messages
|
||||
|
||||
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
|
||||
|
||||
+22
-19
@@ -307,10 +307,7 @@ endif()
|
||||
#
|
||||
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
|
||||
"csrc/quantization/activation_kernels.cu"
|
||||
"csrc/cuda_utils_kernels.cu"
|
||||
"csrc/torch_bindings.cpp")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
@@ -346,9 +343,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
FetchContent_MakeAvailable(cutlass)
|
||||
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/cutlass_extensions/common.cpp")
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${VLLM_EXT_SRC}"
|
||||
CUDA_ARCHS "${CUDA_ARCHS}")
|
||||
@@ -596,8 +590,11 @@ endif()
|
||||
|
||||
if (VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
|
||||
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/custom_quickreduce.cu"
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
)
|
||||
# if ROCM endif
|
||||
endif()
|
||||
@@ -632,6 +629,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/permute_cols.cu"
|
||||
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
|
||||
"csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
|
||||
"csrc/libtorch_stable/pos_encoding_kernels.cu"
|
||||
@@ -639,6 +637,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/layernorm_kernels.cu"
|
||||
"csrc/libtorch_stable/layernorm_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
|
||||
"csrc/libtorch_stable/attention/merge_attn_states.cu"
|
||||
"csrc/libtorch_stable/sampler.cu"
|
||||
"csrc/libtorch_stable/topk.cu"
|
||||
@@ -653,14 +652,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/cuda_utils_kernels.cu"
|
||||
"csrc/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/permute_cols.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
|
||||
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
|
||||
|
||||
@@ -1075,20 +1072,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.10.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.10.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020A000000000000ULL)
|
||||
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.11.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.11.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
# Needed by CUTLASS kernels
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
elseif(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.10.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.10.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020A000000000000ULL)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
|
||||
endif()
|
||||
|
||||
@@ -1306,7 +1308,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set(VLLM_ROCM_HAS_GFX1100 ON)
|
||||
list(APPEND VLLM_ROCM_EXT_SRC
|
||||
"csrc/rocm/q_gemm_rdna3.cu"
|
||||
"csrc/rocm/q_gemm_rdna3_wmma.cu")
|
||||
"csrc/rocm/q_gemm_rdna3_wmma.cu"
|
||||
"csrc/rocm/moe_q_gemm_rdna3.cu")
|
||||
endif()
|
||||
|
||||
define_extension_target(
|
||||
|
||||
@@ -4,6 +4,7 @@ include requirements/cuda.txt
|
||||
include requirements/rocm.txt
|
||||
include requirements/cpu.txt
|
||||
include CMakeLists.txt
|
||||
include tools/build_rust.py
|
||||
|
||||
recursive-include cmake *
|
||||
recursive-include csrc *
|
||||
|
||||
@@ -1,143 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# benchmark the overhead of disaggregated prefill.
|
||||
# methodology:
|
||||
# - send all request to prefill vLLM instance. It will buffer KV cache.
|
||||
# - then send all request to decode instance.
|
||||
# - The TTFT of decode instance is the overhead.
|
||||
|
||||
set -ex
|
||||
|
||||
kill_gpu_processes() {
|
||||
# kill all processes on GPU.
|
||||
pgrep pt_main_thread | xargs -r kill -9
|
||||
pgrep python3 | xargs -r kill -9
|
||||
# vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
|
||||
pgrep VLLM | xargs -r kill -9
|
||||
sleep 10
|
||||
|
||||
# remove vllm config file
|
||||
rm -rf ~/.config/vllm
|
||||
|
||||
# Print the GPU memory usage
|
||||
# so that we know if all GPU processes are killed.
|
||||
gpu_memory_usage=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0)
|
||||
# The memory usage should be 0 MB.
|
||||
echo "GPU 0 Memory Usage: $gpu_memory_usage MB"
|
||||
}
|
||||
|
||||
wait_for_server() {
|
||||
# wait for vllm server to start
|
||||
# return 1 if vllm server crashes
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -s localhost:${port}/v1/completions > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
|
||||
benchmark() {
|
||||
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
|
||||
|
||||
# compare chunked prefill with disaggregated prefill
|
||||
|
||||
results_folder="./results"
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
dataset_name="sonnet"
|
||||
dataset_path="../sonnet_4x.txt"
|
||||
num_prompts=10
|
||||
qps=$1
|
||||
prefix_len=50
|
||||
input_len=2048
|
||||
output_len=$2
|
||||
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
|
||||
--port 8100 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
|
||||
--port 8200 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
|
||||
# let the prefill instance finish prefill
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $model \
|
||||
--dataset-name $dataset_name \
|
||||
--dataset-path $dataset_path \
|
||||
--sonnet-input-len $input_len \
|
||||
--sonnet-output-len "$output_len" \
|
||||
--sonnet-prefix-len $prefix_len \
|
||||
--num-prompts $num_prompts \
|
||||
--port 8100 \
|
||||
--save-result \
|
||||
--result-dir $results_folder \
|
||||
--result-filename disagg_prefill_tp1.json \
|
||||
--request-rate "inf"
|
||||
|
||||
|
||||
# send the request to decode.
|
||||
# The TTFT of this command will be the overhead of disagg prefill impl.
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $model \
|
||||
--dataset-name $dataset_name \
|
||||
--dataset-path $dataset_path \
|
||||
--sonnet-input-len $input_len \
|
||||
--sonnet-output-len "$output_len" \
|
||||
--sonnet-prefix-len $prefix_len \
|
||||
--num-prompts $num_prompts \
|
||||
--port 8200 \
|
||||
--save-result \
|
||||
--result-dir $results_folder \
|
||||
--result-filename disagg_prefill_tp1_overhead.json \
|
||||
--request-rate "$qps"
|
||||
kill_gpu_processes
|
||||
|
||||
}
|
||||
|
||||
|
||||
main() {
|
||||
|
||||
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
|
||||
(which jq) || (apt-get -y install jq)
|
||||
(which socat) || (apt-get -y install socat)
|
||||
|
||||
pip install quart httpx datasets
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
cd ..
|
||||
# create sonnet-4x.txt
|
||||
echo "" > sonnet_4x.txt
|
||||
for _ in {1..4}
|
||||
do
|
||||
cat sonnet.txt >> sonnet_4x.txt
|
||||
done
|
||||
cd disagg_benchmarks
|
||||
|
||||
rm -rf results
|
||||
mkdir results
|
||||
|
||||
default_qps=1
|
||||
default_output_len=1
|
||||
benchmark $default_qps $default_output_len
|
||||
|
||||
}
|
||||
|
||||
|
||||
main "$@"
|
||||
@@ -1,157 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Requirement: 2x GPUs.
|
||||
|
||||
|
||||
# Model: meta-llama/Meta-Llama-3.1-8B-Instruct
|
||||
# Query: 1024 input tokens, 6 output tokens, QPS 2/4/6/8, 100 requests
|
||||
# Resource: 2x GPU
|
||||
# Approaches:
|
||||
# 2. Chunked prefill: 2 vllm instance with tp=4, equivalent to 1 tp=4 instance with QPS 4
|
||||
# 3. Disaggregated prefill: 1 prefilling instance and 1 decoding instance
|
||||
# Prefilling instance: max_output_token=1
|
||||
# Decoding instance: force the input tokens be the same across requests to bypass prefilling
|
||||
|
||||
set -ex
|
||||
|
||||
kill_gpu_processes() {
|
||||
# kill all processes on GPU.
|
||||
pgrep pt_main_thread | xargs -r kill -9
|
||||
pgrep python3 | xargs -r kill -9
|
||||
# vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
|
||||
pgrep VLLM | xargs -r kill -9
|
||||
for port in 8000 8100 8200; do lsof -t -i:$port | xargs -r kill -9; done
|
||||
sleep 1
|
||||
}
|
||||
|
||||
wait_for_server() {
|
||||
# wait for vllm server to start
|
||||
# return 1 if vllm server crashes
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -s localhost:${port}/v1/completions > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
|
||||
launch_chunked_prefill() {
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
# disagg prefill
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
|
||||
--port 8100 \
|
||||
--max-model-len 10000 \
|
||||
--enable-chunked-prefill \
|
||||
--gpu-memory-utilization 0.6 &
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
|
||||
--port 8200 \
|
||||
--max-model-len 10000 \
|
||||
--enable-chunked-prefill \
|
||||
--gpu-memory-utilization 0.6 &
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
python3 round_robin_proxy.py &
|
||||
sleep 1
|
||||
}
|
||||
|
||||
|
||||
launch_disagg_prefill() {
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
# disagg prefill
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
|
||||
--port 8100 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
|
||||
--port 8200 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
python3 disagg_prefill_proxy_server.py &
|
||||
sleep 1
|
||||
}
|
||||
|
||||
|
||||
benchmark() {
|
||||
results_folder="./results"
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
dataset_name="sonnet"
|
||||
dataset_path="../sonnet_4x.txt"
|
||||
num_prompts=100
|
||||
qps=$1
|
||||
prefix_len=50
|
||||
input_len=1024
|
||||
output_len=$2
|
||||
tag=$3
|
||||
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $model \
|
||||
--dataset-name $dataset_name \
|
||||
--dataset-path $dataset_path \
|
||||
--sonnet-input-len $input_len \
|
||||
--sonnet-output-len "$output_len" \
|
||||
--sonnet-prefix-len $prefix_len \
|
||||
--num-prompts $num_prompts \
|
||||
--port 8000 \
|
||||
--save-result \
|
||||
--result-dir $results_folder \
|
||||
--result-filename "$tag"-qps-"$qps".json \
|
||||
--request-rate "$qps"
|
||||
|
||||
sleep 2
|
||||
}
|
||||
|
||||
|
||||
main() {
|
||||
|
||||
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
|
||||
(which jq) || (apt-get -y install jq)
|
||||
(which socat) || (apt-get -y install socat)
|
||||
(which lsof) || (apt-get -y install lsof)
|
||||
|
||||
pip install quart httpx matplotlib aiohttp datasets
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
cd ..
|
||||
# create sonnet-4x.txt so that we can sample 2048 tokens for input
|
||||
echo "" > sonnet_4x.txt
|
||||
for _ in {1..4}
|
||||
do
|
||||
cat sonnet.txt >> sonnet_4x.txt
|
||||
done
|
||||
cd disagg_benchmarks
|
||||
|
||||
rm -rf results
|
||||
mkdir results
|
||||
|
||||
default_output_len=6
|
||||
|
||||
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
|
||||
|
||||
launch_chunked_prefill
|
||||
for qps in 2 4 6 8; do
|
||||
benchmark $qps $default_output_len chunked_prefill
|
||||
done
|
||||
kill_gpu_processes
|
||||
|
||||
launch_disagg_prefill
|
||||
for qps in 2 4 6 8; do
|
||||
benchmark $qps $default_output_len disagg_prefill
|
||||
done
|
||||
kill_gpu_processes
|
||||
|
||||
python3 visualize_benchmark_results.py
|
||||
|
||||
}
|
||||
|
||||
|
||||
main "$@"
|
||||
@@ -1,260 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import aiohttp
|
||||
from quart import Quart, Response, make_response, request
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def parse_args():
|
||||
"""parse command line arguments"""
|
||||
parser = argparse.ArgumentParser(description="vLLM P/D disaggregation proxy server")
|
||||
|
||||
# Add args
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
type=float,
|
||||
default=6 * 60 * 60,
|
||||
help="Timeout for backend service requests in seconds (default: 21600)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8000,
|
||||
help="Port to run the server on (default: 8000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefill-url",
|
||||
type=str,
|
||||
default="http://localhost:8100",
|
||||
help="Prefill service base URL (protocol + host[:port])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--decode-url",
|
||||
type=str,
|
||||
default="http://localhost:8200",
|
||||
help="Decode service base URL (protocol + host[:port])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-host",
|
||||
type=str,
|
||||
default="localhost",
|
||||
help="Hostname or IP used by KV transfer (default: localhost)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefill-kv-port",
|
||||
type=int,
|
||||
default=14579,
|
||||
help="Prefill KV port (default: 14579)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--decode-kv-port",
|
||||
type=int,
|
||||
default=14580,
|
||||
help="Decode KV port (default: 14580)",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
"""parse command line arguments"""
|
||||
args = parse_args()
|
||||
|
||||
# Initialize configuration using command line parameters
|
||||
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
|
||||
PREFILL_SERVICE_URL = args.prefill_url
|
||||
DECODE_SERVICE_URL = args.decode_url
|
||||
PORT = args.port
|
||||
|
||||
PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
|
||||
DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
|
||||
|
||||
logger.info(
|
||||
"Proxy resolved KV addresses -> prefill: %s, decode: %s",
|
||||
PREFILL_KV_ADDR,
|
||||
DECODE_KV_ADDR,
|
||||
)
|
||||
|
||||
app = Quart(__name__)
|
||||
|
||||
# Attach the configuration object to the application instance so helper
|
||||
# coroutines can read the resolved backend URLs and timeouts without using
|
||||
# globals.
|
||||
app.config.update(
|
||||
{
|
||||
"AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
|
||||
"PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
|
||||
"DECODE_SERVICE_URL": DECODE_SERVICE_URL,
|
||||
"PREFILL_KV_ADDR": PREFILL_KV_ADDR,
|
||||
"DECODE_KV_ADDR": DECODE_KV_ADDR,
|
||||
}
|
||||
)
|
||||
|
||||
def _normalize_base_url(url: str) -> str:
|
||||
"""Remove any trailing slash so path joins behave predictably."""
|
||||
return url.rstrip("/")
|
||||
|
||||
def _get_host_port(url: str) -> str:
|
||||
"""Return the hostname:port portion for logging and KV headers."""
|
||||
parsed = urlparse(url)
|
||||
host = parsed.hostname or "localhost"
|
||||
port = parsed.port
|
||||
if port is None:
|
||||
port = 80 if parsed.scheme == "http" else 443
|
||||
return f"{host}:{port}"
|
||||
|
||||
PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
|
||||
DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
|
||||
KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
|
||||
|
||||
def _build_headers(request_id: str) -> dict[str, str]:
|
||||
"""Construct the headers expected by vLLM's P2P disagg connector."""
|
||||
headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
async def _run_prefill(
|
||||
request_path: str,
|
||||
payload: dict,
|
||||
headers: dict[str, str],
|
||||
request_id: str,
|
||||
):
|
||||
url = f"{PREFILL_BASE}{request_path}"
|
||||
start_ts = time.perf_counter()
|
||||
logger.info("[prefill] start request_id=%s url=%s", request_id, url)
|
||||
try:
|
||||
async with (
|
||||
aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
|
||||
session.post(url=url, json=payload, headers=headers) as resp,
|
||||
):
|
||||
if resp.status != 200:
|
||||
error_text = await resp.text()
|
||||
raise RuntimeError(
|
||||
f"Prefill backend error {resp.status}: {error_text}"
|
||||
)
|
||||
await resp.read()
|
||||
logger.info(
|
||||
"[prefill] done request_id=%s status=%s elapsed=%.2fs",
|
||||
request_id,
|
||||
resp.status,
|
||||
time.perf_counter() - start_ts,
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(f"Prefill service timeout at {url}") from exc
|
||||
except aiohttp.ClientError as exc:
|
||||
raise RuntimeError(f"Prefill service unavailable at {url}") from exc
|
||||
|
||||
async def _stream_decode(
|
||||
request_path: str,
|
||||
payload: dict,
|
||||
headers: dict[str, str],
|
||||
request_id: str,
|
||||
):
|
||||
url = f"{DECODE_BASE}{request_path}"
|
||||
# Stream tokens from the decode service once the prefill stage has
|
||||
# materialized KV caches on the target workers.
|
||||
logger.info("[decode] start request_id=%s url=%s", request_id, url)
|
||||
try:
|
||||
async with (
|
||||
aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
|
||||
session.post(url=url, json=payload, headers=headers) as resp,
|
||||
):
|
||||
if resp.status != 200:
|
||||
error_text = await resp.text()
|
||||
logger.error(
|
||||
"Decode backend error %s - %s", resp.status, error_text
|
||||
)
|
||||
err_msg = (
|
||||
'{"error": "Decode backend error ' + str(resp.status) + '"}'
|
||||
)
|
||||
yield err_msg.encode()
|
||||
return
|
||||
logger.info(
|
||||
"[decode] streaming response request_id=%s status=%s",
|
||||
request_id,
|
||||
resp.status,
|
||||
)
|
||||
async for chunk_bytes in resp.content.iter_chunked(1024):
|
||||
yield chunk_bytes
|
||||
logger.info("[decode] finished streaming request_id=%s", request_id)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error("Decode service timeout at %s", url)
|
||||
yield b'{"error": "Decode service timeout"}'
|
||||
except aiohttp.ClientError as exc:
|
||||
logger.error("Decode service error at %s: %s", url, exc)
|
||||
yield b'{"error": "Decode service unavailable"}'
|
||||
|
||||
async def process_request():
|
||||
"""Process a single request through prefill and decode stages"""
|
||||
try:
|
||||
original_request_data = await request.get_json()
|
||||
|
||||
# Create prefill request (max_tokens=1)
|
||||
prefill_request = original_request_data.copy()
|
||||
prefill_request["max_tokens"] = 1
|
||||
if "max_completion_tokens" in prefill_request:
|
||||
prefill_request["max_completion_tokens"] = 1
|
||||
|
||||
# Execute prefill stage
|
||||
# The request id encodes both KV socket addresses so the backend can
|
||||
# shuttle tensors directly via NCCL once the prefill response
|
||||
# completes.
|
||||
request_id = (
|
||||
f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
|
||||
f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
|
||||
)
|
||||
|
||||
headers = _build_headers(request_id)
|
||||
await _run_prefill(request.path, prefill_request, headers, request_id)
|
||||
|
||||
# Execute decode stage and stream response
|
||||
# Pass the unmodified user request so the decode phase can continue
|
||||
# sampling with the already-populated KV cache.
|
||||
generator = _stream_decode(
|
||||
request.path, original_request_data, headers, request_id
|
||||
)
|
||||
response = await make_response(generator)
|
||||
response.timeout = None # Disable timeout for streaming response
|
||||
return response
|
||||
|
||||
except Exception:
|
||||
logger.exception("Error processing request")
|
||||
return Response(
|
||||
response=b'{"error": "Internal server error"}',
|
||||
status=500,
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
@app.route("/v1/completions", methods=["POST"])
|
||||
async def handle_request():
|
||||
"""Handle incoming API requests with concurrency and rate limiting"""
|
||||
try:
|
||||
return await process_request()
|
||||
except asyncio.CancelledError:
|
||||
logger.warning("Request cancelled")
|
||||
return Response(
|
||||
response=b'{"error": "Request cancelled"}',
|
||||
status=503,
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
# Start the Quart server with host can be set to 0.0.0.0
|
||||
app.run(port=PORT)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,63 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
|
||||
class RoundRobinProxy:
|
||||
def __init__(self, target_ports):
|
||||
self.target_ports = target_ports
|
||||
self.port_cycle = itertools.cycle(self.target_ports)
|
||||
|
||||
async def handle_request(self, request):
|
||||
target_port = next(self.port_cycle)
|
||||
target_url = f"http://localhost:{target_port}{request.path_qs}"
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
# Forward the request
|
||||
async with session.request(
|
||||
method=request.method,
|
||||
url=target_url,
|
||||
headers=request.headers,
|
||||
data=request.content,
|
||||
) as response:
|
||||
# Start sending the response
|
||||
resp = web.StreamResponse(
|
||||
status=response.status, headers=response.headers
|
||||
)
|
||||
await resp.prepare(request)
|
||||
|
||||
# Stream the response content
|
||||
async for chunk in response.content.iter_any():
|
||||
await resp.write(chunk)
|
||||
|
||||
await resp.write_eof()
|
||||
return resp
|
||||
|
||||
except Exception as e:
|
||||
return web.Response(text=f"Error: {str(e)}", status=500)
|
||||
|
||||
|
||||
async def main():
|
||||
proxy = RoundRobinProxy([8100, 8200])
|
||||
app = web.Application()
|
||||
app.router.add_route("*", "/{path:.*}", proxy.handle_request)
|
||||
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "localhost", 8000)
|
||||
await site.start()
|
||||
|
||||
print("Proxy server started on http://localhost:8000")
|
||||
|
||||
# Keep the server running
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -1,47 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import pandas as pd
|
||||
|
||||
if __name__ == "__main__":
|
||||
data = []
|
||||
for name in ["disagg_prefill", "chunked_prefill"]:
|
||||
for qps in [2, 4, 6, 8]:
|
||||
with open(f"results/{name}-qps-{qps}.json") as f:
|
||||
x = json.load(f)
|
||||
x["name"] = name
|
||||
x["qps"] = qps
|
||||
data.append(x)
|
||||
|
||||
df = pd.DataFrame.from_dict(data)
|
||||
dis_df = df[df["name"] == "disagg_prefill"]
|
||||
chu_df = df[df["name"] == "chunked_prefill"]
|
||||
|
||||
plt.style.use("bmh")
|
||||
plt.rcParams["font.size"] = 20
|
||||
|
||||
for key in [
|
||||
"mean_ttft_ms",
|
||||
"median_ttft_ms",
|
||||
"p99_ttft_ms",
|
||||
"mean_itl_ms",
|
||||
"median_itl_ms",
|
||||
"p99_itl_ms",
|
||||
]:
|
||||
fig, ax = plt.subplots(figsize=(11, 7))
|
||||
plt.plot(
|
||||
dis_df["qps"], dis_df[key], label="disagg_prefill", marker="o", linewidth=4
|
||||
)
|
||||
plt.plot(
|
||||
chu_df["qps"], chu_df[key], label="chunked_prefill", marker="o", linewidth=4
|
||||
)
|
||||
ax.legend()
|
||||
|
||||
ax.set_xlabel("QPS")
|
||||
ax.set_ylabel(key)
|
||||
ax.set_ylim(bottom=0)
|
||||
fig.savefig(f"results/{key}.png")
|
||||
plt.close(fig)
|
||||
@@ -250,7 +250,7 @@ def benchmark_config(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=topk,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size_per_partition=shard_intermediate_size,
|
||||
intermediate_size=shard_intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
activation=MoEActivation.SILU,
|
||||
@@ -271,7 +271,6 @@ def benchmark_config(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
),
|
||||
inplace=not disable_inplace(),
|
||||
)
|
||||
|
||||
with override_config(config):
|
||||
@@ -279,7 +278,6 @@ def benchmark_config(
|
||||
x, input_gating, topk, renormalize=not use_deep_gemm
|
||||
)
|
||||
|
||||
inplace = not disable_inplace()
|
||||
if use_deep_gemm:
|
||||
return deep_gemm_experts.apply(
|
||||
x,
|
||||
@@ -298,7 +296,6 @@ def benchmark_config(
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
inplace=inplace,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -65,6 +65,32 @@ class RequestArgs(NamedTuple):
|
||||
limit_min_tokens: int # Use negative value for no limit
|
||||
limit_max_tokens: int # Use negative value for no limit
|
||||
timeout_sec: int
|
||||
send_conversation_id: bool
|
||||
headers: dict[str, str]
|
||||
|
||||
|
||||
def parse_custom_header(header: str) -> tuple[str, str]:
|
||||
separators = (":", "=")
|
||||
for separator in separators:
|
||||
if separator in header:
|
||||
key, value = header.split(separator, 1)
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
if key:
|
||||
return key, value
|
||||
break
|
||||
raise argparse.ArgumentTypeError(
|
||||
"Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
|
||||
)
|
||||
|
||||
|
||||
def build_request_headers(
|
||||
api_key: str | None, custom_headers: list[tuple[str, str]] | None
|
||||
) -> dict[str, str]:
|
||||
headers = dict(custom_headers or [])
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
|
||||
class BenchmarkArgs(NamedTuple):
|
||||
@@ -218,12 +244,11 @@ async def send_request(
|
||||
max_tokens: int | None = None,
|
||||
timeout_sec: int = 120,
|
||||
conversation_id: str | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> ServerResponse:
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"seed": 0,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
if conversation_id is not None:
|
||||
@@ -233,13 +258,17 @@ async def send_request(
|
||||
payload["stream"] = True
|
||||
payload["stream_options"] = {"include_usage": False}
|
||||
|
||||
if min_tokens is not None:
|
||||
payload["min_tokens"] = min_tokens
|
||||
# if min_tokens is not None:
|
||||
# payload["min_tokens"] = min_tokens
|
||||
|
||||
if max_tokens is not None:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
headers = {"Content-Type": "application/json"}
|
||||
request_headers = {"Content-Type": "application/json"}
|
||||
if conversation_id is not None:
|
||||
request_headers["X-Session-ID"] = str(conversation_id)
|
||||
if headers is not None:
|
||||
request_headers.update(headers)
|
||||
|
||||
# Calculate the timeout for the request
|
||||
if max_tokens is not None:
|
||||
@@ -265,7 +294,7 @@ async def send_request(
|
||||
most_recent_timestamp: int = start_time
|
||||
|
||||
async with session.post(
|
||||
url=chat_url, json=payload, headers=headers, timeout=timeout
|
||||
url=chat_url, json=payload, headers=request_headers, timeout=timeout
|
||||
) as response:
|
||||
http_status = HTTPStatus(response.status)
|
||||
if http_status == HTTPStatus.OK:
|
||||
@@ -317,6 +346,8 @@ async def send_request(
|
||||
latency = time.perf_counter_ns() - start_time
|
||||
|
||||
if ttft is None:
|
||||
if stream:
|
||||
valid_response = False
|
||||
# The response was a single chunk
|
||||
ttft = latency
|
||||
|
||||
@@ -423,7 +454,8 @@ async def send_turn(
|
||||
min_tokens,
|
||||
max_tokens,
|
||||
req_args.timeout_sec,
|
||||
conversation_id=conv_id,
|
||||
conversation_id=conv_id if req_args.send_conversation_id else None,
|
||||
headers=req_args.headers,
|
||||
)
|
||||
|
||||
if response.valid is False:
|
||||
@@ -872,6 +904,7 @@ def get_client_config(
|
||||
# Arguments for API requests
|
||||
chat_url = f"{args.url}/v1/chat/completions"
|
||||
model_name = args.served_model_name if args.served_model_name else args.model
|
||||
headers = build_request_headers(args.api_key, args.header)
|
||||
|
||||
req_args = RequestArgs(
|
||||
chat_url=chat_url,
|
||||
@@ -880,6 +913,8 @@ def get_client_config(
|
||||
limit_min_tokens=args.limit_min_tokens,
|
||||
limit_max_tokens=args.limit_max_tokens,
|
||||
timeout_sec=args.request_timeout_sec,
|
||||
send_conversation_id=args.send_conversation_id,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
return client_args, req_args
|
||||
@@ -1245,19 +1280,19 @@ def process_statistics(
|
||||
)
|
||||
|
||||
|
||||
async def get_server_info(url: str) -> None:
|
||||
async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
|
||||
logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Get server version (not mandatory, "version" endpoint may not exist)
|
||||
url_version = f"{url}/version"
|
||||
async with session.get(url_version) as response:
|
||||
async with session.get(url_version, headers=headers) as response:
|
||||
if HTTPStatus(response.status) == HTTPStatus.OK:
|
||||
text = await response.text()
|
||||
logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
|
||||
|
||||
# Get available models
|
||||
url_models = f"{url}/v1/models"
|
||||
async with session.get(url_models) as response:
|
||||
async with session.get(url_models, headers=headers) as response:
|
||||
if HTTPStatus(response.status) == HTTPStatus.OK:
|
||||
text = await response.text()
|
||||
logger.info(f"{Color.BLUE}Models:{Color.RESET}")
|
||||
@@ -1323,6 +1358,22 @@ async def main() -> None:
|
||||
help="Base URL for the LLM API server",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--api-key",
|
||||
type=str,
|
||||
default=None,
|
||||
help="API key to send as an Authorization bearer token",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--header",
|
||||
action="append",
|
||||
type=parse_custom_header,
|
||||
default=None,
|
||||
metavar="KEY=VALUE",
|
||||
help="Custom request header. Can be specified multiple times. "
|
||||
"Accepts 'Header-Name: value' or 'Header-Name=value'.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--num-clients",
|
||||
@@ -1437,6 +1488,22 @@ async def main() -> None:
|
||||
help="Disable stream/streaming mode (set 'stream' to False in the API request)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--send-conversation-id",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help=(
|
||||
"Inject a `conversation_id` field into each Chat Completions "
|
||||
"payload. This is a non-standard OpenAI extension consumed by "
|
||||
"vLLM's disaggregated multi-turn proxy "
|
||||
"(examples/disaggregated/disaggregated_serving/"
|
||||
"disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. "
|
||||
"Leave disabled (default) when targeting strict "
|
||||
"OpenAI-compatible endpoints; enable when benchmarking the "
|
||||
"disaggregated proxy."
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-e",
|
||||
"--excel-output",
|
||||
@@ -1525,7 +1592,8 @@ async def main() -> None:
|
||||
args.model, trust_remote_code=args.trust_remote_code
|
||||
)
|
||||
|
||||
await get_server_info(args.url)
|
||||
headers = build_request_headers(args.api_key, args.header)
|
||||
await get_server_info(args.url, headers=headers)
|
||||
|
||||
# Load the input file (either conversations of configuration file)
|
||||
logger.info(f"Reading input file: {args.input_file}")
|
||||
|
||||
+4
-15
@@ -1,5 +1,5 @@
|
||||
#!/bin/bash
|
||||
# Build the vllm-rs Rust frontend binary and install it into the vllm package.
|
||||
# Build vLLM Rust artifacts and install them into the vllm package.
|
||||
# Usage: ./build_rust.sh [--debug]
|
||||
#
|
||||
# By default builds in release mode. Pass --debug for faster compile times
|
||||
@@ -8,8 +8,6 @@
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
|
||||
RUST_DIR="$REPO_ROOT/rust"
|
||||
TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
|
||||
|
||||
# Read the required toolchain from rust-toolchain.toml.
|
||||
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
|
||||
@@ -27,18 +25,9 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
|
||||
fi
|
||||
|
||||
if [[ "${1:-}" == "--debug" ]]; then
|
||||
PROFILE_ARGS=()
|
||||
PROFILE_DIR="debug"
|
||||
PROFILE_ARG="--debug"
|
||||
else
|
||||
PROFILE_ARGS=(--release)
|
||||
PROFILE_DIR="release"
|
||||
PROFILE_ARG="--release"
|
||||
fi
|
||||
|
||||
cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
|
||||
--manifest-path "$RUST_DIR/Cargo.toml" \
|
||||
--bin vllm-rs \
|
||||
--features native-tls-vendored
|
||||
|
||||
mkdir -p "$(dirname "$TARGET_PATH")"
|
||||
cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
|
||||
echo "Installed vllm-rs to $TARGET_PATH"
|
||||
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#include <torch/all.h>
|
||||
#include <torch/cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
static const char* PYARGS_PARSE = "KKKK";
|
||||
#else
|
||||
#include <cstdlib>
|
||||
#include <cstdint>
|
||||
#include <cerrno>
|
||||
#include <climits>
|
||||
|
||||
@@ -46,6 +47,29 @@ static inline unsigned long long my_min(unsigned long long a,
|
||||
return a < b ? a : b;
|
||||
}
|
||||
|
||||
static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,
|
||||
size_t alignment) {
|
||||
CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0);
|
||||
if (status == CUresult(0) || alignment == 0) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Some ROCm stacks can report OOM while reserving VA with an explicit
|
||||
// alignment even when physical VRAM is free. Let HIP choose the default
|
||||
// alignment, then verify that the returned address still satisfies the
|
||||
// requested alignment before accepting it.
|
||||
status = cuMemAddressReserve(d_mem, size, 0, 0, 0);
|
||||
if (status != CUresult(0)) {
|
||||
return status;
|
||||
}
|
||||
if (((std::uintptr_t)(*d_mem) % alignment) == 0) {
|
||||
return status;
|
||||
}
|
||||
|
||||
(void)cuMemAddressFree(*d_mem, size);
|
||||
return hipErrorNotSupported;
|
||||
}
|
||||
|
||||
static const char* PYARGS_PARSE = "KKKO";
|
||||
#endif
|
||||
|
||||
@@ -325,7 +349,7 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
|
||||
return nullptr;
|
||||
}
|
||||
#else
|
||||
CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, granularity, 0, 0));
|
||||
CUDA_CHECK(reserve_rocm_address(&d_mem, alignedSize, granularity));
|
||||
if (error_code != 0) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
#ifndef CONCAT_MLA_Q_CUH_
|
||||
#define CONCAT_MLA_Q_CUH_
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include "cuda_vec_utils.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include <torch/csrc/stable/accelerator.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
#include <torch/csrc/stable/device.h>
|
||||
#include <torch/csrc/stable/c/shim.h>
|
||||
#include <torch/headeronly/version.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <array>
|
||||
#include <optional>
|
||||
|
||||
// This function assumes that `cpu_tensor` is a CPU tensor,
|
||||
// and that UVA (Unified Virtual Addressing) is enabled.
|
||||
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
|
||||
torch::stable::Tensor& cpu_tensor) {
|
||||
STD_TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
|
||||
|
||||
const auto dtype = cpu_tensor.scalar_type();
|
||||
const auto layout = cpu_tensor.layout();
|
||||
const torch::stable::Device cuda_dev(torch::headeronly::DeviceType::CUDA);
|
||||
|
||||
// handle empty tensor
|
||||
if (cpu_tensor.numel() == 0) {
|
||||
return torch::stable::empty(cpu_tensor.sizes(), dtype, layout, cuda_dev);
|
||||
}
|
||||
|
||||
std::array<StableIValue, 2> is_pinned_stack{
|
||||
torch::stable::detail::from(cpu_tensor),
|
||||
torch::stable::detail::from(std::nullopt)};
|
||||
TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
|
||||
"aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION));
|
||||
if (torch::stable::detail::to<bool>(is_pinned_stack[0])) {
|
||||
// If CPU tensor is pinned, directly get the device pointer.
|
||||
void* host_ptr = const_cast<void*>(cpu_tensor.mutable_data_ptr());
|
||||
void* device_ptr = nullptr;
|
||||
cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
STD_TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed: ",
|
||||
cudaGetErrorString(err));
|
||||
|
||||
return torch::stable::from_blob(
|
||||
device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), cuda_dev, dtype,
|
||||
[base = cpu_tensor](void*) {}); // keep cpu tensor alive
|
||||
}
|
||||
|
||||
// If CPU tensor is not pinned, allocate a new pinned memory buffer.
|
||||
torch::stable::Tensor contiguous_cpu = torch::stable::contiguous(cpu_tensor);
|
||||
size_t nbytes = contiguous_cpu.numel() * contiguous_cpu.element_size();
|
||||
|
||||
void* host_ptr = nullptr;
|
||||
cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
|
||||
if (err != cudaSuccess) {
|
||||
STD_TORCH_CHECK(false, "cudaHostAlloc failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
err = cudaMemcpy(host_ptr, contiguous_cpu.const_data_ptr(), nbytes,
|
||||
cudaMemcpyDefault);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
STD_TORCH_CHECK(false, "cudaMemcpy failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
void* device_ptr = nullptr;
|
||||
err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
STD_TORCH_CHECK(
|
||||
false, "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
|
||||
|
||||
return torch::stable::from_blob(device_ptr, contiguous_cpu.sizes(),
|
||||
contiguous_cpu.strides(), cuda_dev,
|
||||
contiguous_cpu.scalar_type(), deleter);
|
||||
}
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "common.hpp"
|
||||
|
||||
int32_t get_sm_version_num() {
|
||||
int32_t major_capability, minor_capability;
|
||||
@@ -102,6 +102,35 @@ constexpr float NUM_TOKEN_CUTOFF = 1024;
|
||||
constexpr int kNumLanes = 32;
|
||||
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
|
||||
|
||||
// Pack this lane's 16 fp32 elements into per-tensor E4M3 FP8 (one uint4 = 16
|
||||
// B), scaling by `scale` (a reciprocal scale) and saturating to ±448. Used by
|
||||
// the FlashInfer full-cache path for both the Q and KV stores.
|
||||
__device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
|
||||
float const scale) {
|
||||
#ifndef USE_ROCM
|
||||
uint4 out;
|
||||
auto* out2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane / 2; i++) {
|
||||
float2 scaled =
|
||||
make_float2(values[2 * i] * scale, values[2 * i + 1] * scale);
|
||||
scaled.x = fminf(fmaxf(scaled.x, -kFp8Max), kFp8Max);
|
||||
scaled.y = fminf(fmaxf(scaled.y, -kFp8Max), kFp8Max);
|
||||
out2[i] = __nv_cvt_float2_to_fp8x2(scaled, __NV_SATFINITE, __NV_E4M3);
|
||||
}
|
||||
return out;
|
||||
#else
|
||||
uint8_t out_bytes[kElemsPerLane];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float scaled = values[i] * scale;
|
||||
scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
|
||||
out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
|
||||
}
|
||||
return *reinterpret_cast<uint4 const*>(out_bytes);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Small inline helpers
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
@@ -649,6 +678,257 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
|
||||
#undef DISPATCH
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// FlashInfer full-cache kernel
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
//
|
||||
// Sibling to the FlashMLA kernel above, used by the FlashInfer V4 sparse-MLA
|
||||
// backend. Differences from the legacy path:
|
||||
// * No Q head padding — output Q layout matches the input num_heads_q.
|
||||
// * KV is written as a *contiguous* 512-wide row per token (token-strided),
|
||||
// not the legacy UE8M0 paged layout with a separate scale tail.
|
||||
// * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global
|
||||
// scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags.
|
||||
//
|
||||
// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / warps).
|
||||
// Each warp handles one (token, slot): slot < num_heads_q → Q, slot ==
|
||||
// num_heads_q → KV.
|
||||
template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
|
||||
__global__ void fusedDeepseekV4FullCacheKernel(
|
||||
scalar_t_in* __restrict__ q_inout, // [N, H, 512], in place (bf16)
|
||||
uint8_t* __restrict__ q_fp8_out, // [N, H, 512] fp8, optional
|
||||
int64_t const q_fp8_stride0, // elements (fp8 == bytes)
|
||||
int64_t const q_fp8_stride1, // elements (fp8 == bytes)
|
||||
scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
|
||||
uint8_t* __restrict__ k_cache, // contiguous bf16 or fp8 cache
|
||||
int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
|
||||
int64_t const* __restrict__ position_ids, // [N] i64
|
||||
float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
|
||||
float const* __restrict__ fp8_scale_ptr, // scalar, KV fp8 only
|
||||
float const* __restrict__ q_fp8_scale_inv, // scalar, Q fp8 only
|
||||
float const eps,
|
||||
int const num_tokens_full, // = q.size(0) = kv.size(0)
|
||||
int const num_tokens_insert, // = slot_mapping.size(0)
|
||||
int const num_heads_q, // H (no padding)
|
||||
int const cache_block_size, // tokens per cache block
|
||||
int64_t const kv_block_stride, // bytes per cache block
|
||||
int64_t const kv_token_stride) { // bytes per cache token
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
using Converter = vllm::_typeConvert<scalar_t_in>;
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const warpId = threadIdx.x / 32;
|
||||
int const laneId = threadIdx.x % 32;
|
||||
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
|
||||
|
||||
int const slotsPerToken = num_heads_q + 1;
|
||||
int const tokenIdx = globalWarpIdx / slotsPerToken;
|
||||
int const slotIdx = globalWarpIdx % slotsPerToken;
|
||||
if (tokenIdx >= num_tokens_full) return;
|
||||
bool const isKV = (slotIdx == num_heads_q);
|
||||
// KV branch: skip DP-padded tokens (no slot reserved for them).
|
||||
if (isKV && tokenIdx >= num_tokens_insert) return;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
|
||||
scalar_t_in const* src_ptr;
|
||||
if (isKV) {
|
||||
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
|
||||
} else {
|
||||
src_ptr = q_inout +
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
|
||||
kHeadDim +
|
||||
dim_base;
|
||||
}
|
||||
uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
|
||||
uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
|
||||
|
||||
// ── Decode bf16 → 16 fp32 registers ───────────────────────────────────
|
||||
float elements[kElemsPerLane];
|
||||
{
|
||||
auto const* p0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
|
||||
auto const* p1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p0[i]);
|
||||
elements[2 * i] = f2.x;
|
||||
elements[2 * i + 1] = f2.y;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p1[i]);
|
||||
elements[8 + 2 * i] = f2.x;
|
||||
elements[8 + 2 * i + 1] = f2.y;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Q branch: RMSNorm (no weight) ─────────────────────────────────────
|
||||
if (!isKV) {
|
||||
float sumOfSquares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
sumOfSquares += elements[i] * elements[i];
|
||||
}
|
||||
sumOfSquares = warpSum<float>(sumOfSquares);
|
||||
float const rms_rcp =
|
||||
rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
elements[i] = elements[i] * rms_rcp;
|
||||
}
|
||||
}
|
||||
|
||||
// ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────
|
||||
bool const is_rope_lane = dim_base >= kNopeDim;
|
||||
if (is_rope_lane) {
|
||||
int64_t const pos = position_ids[tokenIdx];
|
||||
constexpr int kHalfRope = kRopeDim / 2;
|
||||
float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
|
||||
float const* sin_ptr = cos_ptr + kHalfRope;
|
||||
int const rope_local_base = dim_base - kNopeDim;
|
||||
int const half_base = rope_local_base >> 1;
|
||||
float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
|
||||
float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
|
||||
float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
|
||||
float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
|
||||
float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
|
||||
float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
|
||||
#pragma unroll
|
||||
for (int p = 0; p < kElemsPerLane / 2; p++) {
|
||||
float const x_even = elements[2 * p];
|
||||
float const x_odd = elements[2 * p + 1];
|
||||
elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
|
||||
elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
|
||||
}
|
||||
}
|
||||
|
||||
// ── Store ─────────────────────────────────────────────────────────────
|
||||
if (!isKV) {
|
||||
if constexpr (STORE_Q_FP8) {
|
||||
float const scale_inv = VLLM_LDG(q_fp8_scale_inv);
|
||||
uint4 const out = packFp8E4M3x16(elements, scale_inv);
|
||||
uint8_t* dst = q_fp8_out +
|
||||
static_cast<int64_t>(tokenIdx) * q_fp8_stride0 +
|
||||
static_cast<int64_t>(slotIdx) * q_fp8_stride1 + dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out;
|
||||
} else {
|
||||
uint4 out0, out1;
|
||||
auto* po0 = reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
auto* po1 = reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
scalar_t_in* dst =
|
||||
q_inout +
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
|
||||
dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out0;
|
||||
*reinterpret_cast<uint4*>(dst + 8) = out1;
|
||||
}
|
||||
} else {
|
||||
int64_t const slot_id = slot_mapping[tokenIdx];
|
||||
if (slot_id >= 0) {
|
||||
int64_t const block_idx = slot_id / cache_block_size;
|
||||
int64_t const pos_in_block = slot_id % cache_block_size;
|
||||
uint8_t* cache_row =
|
||||
k_cache + block_idx * kv_block_stride + pos_in_block * kv_token_stride;
|
||||
if constexpr (STORE_KV_FP8) {
|
||||
float const inv_scale = 1.0f / VLLM_LDG(fp8_scale_ptr);
|
||||
uint4 const out = packFp8E4M3x16(elements, inv_scale);
|
||||
*reinterpret_cast<uint4*>(cache_row + dim_base) = out;
|
||||
} else {
|
||||
uint4 out0, out1;
|
||||
auto* po0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
auto* po1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
scalar_t_in* dst = reinterpret_cast<scalar_t_in*>(cache_row) + dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out0;
|
||||
*reinterpret_cast<uint4*>(dst + 8) = out1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Configure + launch helper shared by the bf16 and fp8 full-cache launchers.
|
||||
template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
|
||||
static void launchFullCacheKernel(
|
||||
scalar_t_in* q_inout, uint8_t* q_fp8_out, int64_t q_fp8_stride0,
|
||||
int64_t q_fp8_stride1, scalar_t_in const* kv_in, uint8_t* k_cache,
|
||||
int64_t const* slot_mapping, int64_t const* position_ids,
|
||||
float const* cos_sin_cache, float const* fp8_scale,
|
||||
float const* q_fp8_scale_inv, float const eps, int const num_tokens_full,
|
||||
int const num_tokens_insert, int const num_heads_q,
|
||||
int const cache_block_size, int64_t const kv_block_stride,
|
||||
int64_t const kv_token_stride, char const* op_name, cudaStream_t stream) {
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int kWarpsPerBlock = kBlockSize / 32;
|
||||
int64_t const total_warps =
|
||||
static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
|
||||
int const grid =
|
||||
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
auto* kernel =
|
||||
fusedDeepseekV4FullCacheKernel<scalar_t_in, STORE_Q_FP8, STORE_KV_FP8>;
|
||||
#ifndef USE_ROCM
|
||||
static int const sm_version = getSMVersion();
|
||||
STD_TORCH_CHECK(sm_version >= 80, op_name,
|
||||
" requires sm_80+ (Ampere or newer); got sm_", sm_version);
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = dim3(grid);
|
||||
config.blockDim = dim3(kBlockSize);
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.attrs = attrs;
|
||||
config.numAttrs = (sm_version >= 90) ? 1 : 0;
|
||||
cudaLaunchKernelEx(&config, kernel, q_inout, q_fp8_out, q_fp8_stride0,
|
||||
q_fp8_stride1, kv_in, k_cache, slot_mapping, position_ids,
|
||||
cos_sin_cache, fp8_scale, q_fp8_scale_inv, eps,
|
||||
num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
cache_block_size, kv_block_stride, kv_token_stride);
|
||||
#else
|
||||
kernel<<<grid, kBlockSize, 0, stream>>>(
|
||||
q_inout, q_fp8_out, q_fp8_stride0, q_fp8_stride1, kv_in, k_cache,
|
||||
slot_mapping, position_ids, cos_sin_cache, fp8_scale, q_fp8_scale_inv,
|
||||
eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
|
||||
kv_block_stride, kv_token_stride);
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace deepseek_v4_fused_ops
|
||||
} // namespace vllm
|
||||
|
||||
@@ -735,3 +1015,167 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
});
|
||||
return q_out;
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// FlashInfer full-cache torch ops
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
|
||||
torch::stable::Tensor& q, // [N, H, 512] bf16, in place
|
||||
torch::stable::Tensor const& kv, // [N, 512] bf16, read-only
|
||||
torch::stable::Tensor& k_cache, // [num_blocks, bs, 512] bf16
|
||||
torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
|
||||
torch::stable::Tensor const& position_ids, // [N] int64
|
||||
torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
|
||||
double eps, int64_t cache_block_size) {
|
||||
using torch::headeronly::ScalarType;
|
||||
STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
|
||||
"q must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
|
||||
"kv must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
|
||||
STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
|
||||
slot_mapping.scalar_type() == ScalarType::Long,
|
||||
"slot_mapping must be int64 CUDA");
|
||||
STD_TORCH_CHECK(position_ids.device().is_cuda() &&
|
||||
position_ids.scalar_type() == ScalarType::Long,
|
||||
"position_ids must be int64 CUDA");
|
||||
STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
|
||||
cos_sin_cache.scalar_type() == ScalarType::Float &&
|
||||
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
|
||||
"cos_sin_cache shape [max_pos, 64] float32");
|
||||
STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
|
||||
STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
|
||||
STD_TORCH_CHECK(q.scalar_type() == ScalarType::BFloat16 &&
|
||||
kv.scalar_type() == ScalarType::BFloat16,
|
||||
"q and kv must be bfloat16");
|
||||
STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
|
||||
k_cache.size(2) == 512 && k_cache.stride(2) == 1,
|
||||
"k_cache shape [num_blocks, cache_block_size, 512] contiguous");
|
||||
STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::BFloat16,
|
||||
"k_cache must be bfloat16");
|
||||
|
||||
int const num_tokens_full = static_cast<int>(q.size(0));
|
||||
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
|
||||
STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
|
||||
static_cast<int>(position_ids.size(0)) == num_tokens_full,
|
||||
"q/kv/position_ids row counts must match");
|
||||
STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
|
||||
"slot_mapping must not exceed q row count");
|
||||
int const num_heads_q = static_cast<int>(q.size(1));
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
q.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
|
||||
|
||||
// bf16 cache: 2 bytes/element -> byte strides for the uint8-addressed kernel.
|
||||
int64_t const kv_block_stride = k_cache.stride(0) * 2;
|
||||
int64_t const kv_token_stride = k_cache.stride(1) * 2;
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
q.scalar_type(),
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", [&] {
|
||||
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false,
|
||||
false>(
|
||||
reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), nullptr, 0, 0,
|
||||
reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
|
||||
reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
|
||||
slot_mapping.const_data_ptr<int64_t>(),
|
||||
position_ids.const_data_ptr<int64_t>(),
|
||||
cos_sin_cache.const_data_ptr<float>(), nullptr, nullptr,
|
||||
static_cast<float>(eps), num_tokens_full, num_tokens_insert,
|
||||
num_heads_q, static_cast<int>(cache_block_size), kv_block_stride,
|
||||
kv_token_stride,
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
|
||||
torch::stable::Tensor const& q, // [N, H, 512] bf16, read-only
|
||||
torch::stable::Tensor const& kv, // [N, 512] bf16, read-only
|
||||
torch::stable::Tensor& q_fp8, // [N, H, 512] fp8 e4m3
|
||||
torch::stable::Tensor& k_cache, // [num_blocks, bs, 512] fp8
|
||||
torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
|
||||
torch::stable::Tensor const& position_ids, // [N] int64
|
||||
torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
|
||||
torch::stable::Tensor const& fp8_scale, // scalar float32 (KV scale)
|
||||
torch::stable::Tensor const& q_fp8_scale_inv, // scalar float32 (1 / Q scale)
|
||||
double eps, int64_t cache_block_size) {
|
||||
using torch::headeronly::ScalarType;
|
||||
STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
|
||||
"q must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
|
||||
"kv must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(q_fp8.device().is_cuda() && q_fp8.is_contiguous() &&
|
||||
q_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
|
||||
q_fp8.dim() == 3 && q_fp8.size(0) == q.size(0) &&
|
||||
q_fp8.size(1) == q.size(1) && q_fp8.size(2) == q.size(2),
|
||||
"q_fp8 must be a contiguous float8_e4m3fn tensor matching q");
|
||||
STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
|
||||
STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
|
||||
slot_mapping.scalar_type() == ScalarType::Long,
|
||||
"slot_mapping must be int64 CUDA");
|
||||
STD_TORCH_CHECK(position_ids.device().is_cuda() &&
|
||||
position_ids.scalar_type() == ScalarType::Long,
|
||||
"position_ids must be int64 CUDA");
|
||||
STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
|
||||
cos_sin_cache.scalar_type() == ScalarType::Float &&
|
||||
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
|
||||
"cos_sin_cache shape [max_pos, 64] float32");
|
||||
STD_TORCH_CHECK(fp8_scale.device().is_cuda() &&
|
||||
fp8_scale.scalar_type() == ScalarType::Float &&
|
||||
fp8_scale.size(0) == 1,
|
||||
"fp8_scale must be a scalar float32 CUDA tensor");
|
||||
STD_TORCH_CHECK(q_fp8_scale_inv.device().is_cuda() &&
|
||||
q_fp8_scale_inv.scalar_type() == ScalarType::Float &&
|
||||
q_fp8_scale_inv.size(0) == 1,
|
||||
"q_fp8_scale_inv must be a scalar float32 CUDA tensor");
|
||||
STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
|
||||
STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
|
||||
STD_TORCH_CHECK(q.scalar_type() == kv.scalar_type(),
|
||||
"q and kv dtype must match");
|
||||
STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
|
||||
k_cache.size(2) == 512 && k_cache.stride(2) == 1,
|
||||
"k_cache shape [num_blocks, cache_block_size, 512] contiguous");
|
||||
STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
|
||||
"k_cache must be float8_e4m3fn");
|
||||
|
||||
int const num_tokens_full = static_cast<int>(q.size(0));
|
||||
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
|
||||
STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
|
||||
static_cast<int>(position_ids.size(0)) == num_tokens_full,
|
||||
"q/kv/position_ids row counts must match");
|
||||
STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
|
||||
"slot_mapping must not exceed q row count");
|
||||
int const num_heads_q = static_cast<int>(q.size(1));
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
q.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
q.scalar_type(),
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", [&] {
|
||||
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, true,
|
||||
true>(
|
||||
// q is read-only in the fp8 path (the kernel writes q_fp8); the
|
||||
// launcher signature is non-const, so cast away const on the ptr.
|
||||
reinterpret_cast<scalar_t*>(
|
||||
const_cast<void*>(q.const_data_ptr())),
|
||||
reinterpret_cast<uint8_t*>(q_fp8.mutable_data_ptr()),
|
||||
q_fp8.stride(0), q_fp8.stride(1),
|
||||
reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
|
||||
reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
|
||||
slot_mapping.const_data_ptr<int64_t>(),
|
||||
position_ids.const_data_ptr<int64_t>(),
|
||||
cos_sin_cache.const_data_ptr<float>(),
|
||||
fp8_scale.const_data_ptr<float>(),
|
||||
q_fp8_scale_inv.const_data_ptr<float>(), static_cast<float>(eps),
|
||||
num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
static_cast<int>(cache_block_size),
|
||||
// fp8 cache: 1 byte/element -> stride already in bytes.
|
||||
k_cache.stride(0), k_cache.stride(1),
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -24,10 +24,10 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
|
||||
int64_t group_size, double eps, double int8_min,
|
||||
double int8_max);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
|
||||
torch::stable::Tensor const& perm);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
|
||||
@@ -162,6 +162,11 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
|
||||
// AllSpark ops: declarations are in the source files
|
||||
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
|
||||
|
||||
// TODO: Move this out once ROCm upgrade their torch to 2.11.
|
||||
// CPU tensor -> CUDA UVA view (shared CUDA)
|
||||
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
|
||||
torch::stable::Tensor& cpu_tensor);
|
||||
|
||||
#endif
|
||||
|
||||
// Attention kernels (shared CUDA/ROCm)
|
||||
@@ -215,6 +220,13 @@ void rms_norm_per_block_quant(torch::stable::Tensor& out,
|
||||
std::optional<torch::stable::Tensor> residual,
|
||||
int64_t group_size, bool is_scale_transposed);
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor& scales,
|
||||
int64_t group_size,
|
||||
std::optional<torch::stable::Tensor> scale_ub,
|
||||
bool is_scale_transposed);
|
||||
|
||||
// Positional encoding kernels (shared CUDA/ROCm)
|
||||
void rotary_embedding(torch::stable::Tensor& positions,
|
||||
torch::stable::Tensor& query,
|
||||
@@ -238,6 +250,23 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
|
||||
double eps, int64_t cache_block_size);
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
|
||||
torch::stable::Tensor& q, torch::stable::Tensor const& kv,
|
||||
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
|
||||
torch::stable::Tensor const& position_ids,
|
||||
torch::stable::Tensor const& cos_sin_cache, double eps,
|
||||
int64_t cache_block_size);
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
|
||||
torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
|
||||
torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
|
||||
torch::stable::Tensor const& slot_mapping,
|
||||
torch::stable::Tensor const& position_ids,
|
||||
torch::stable::Tensor const& cos_sin_cache,
|
||||
torch::stable::Tensor const& fp8_scale,
|
||||
torch::stable::Tensor const& q_fp8_scale_inv, double eps,
|
||||
int64_t cache_block_size);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
torch::stable::Tensor minimax_allreduce_rms(
|
||||
torch::stable::Tensor const& input,
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
#include "cutlass_extensions/torch_utils.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "get_group_starts.cuh"
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
|
||||
@@ -21,7 +21,7 @@
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
#include "cutlass/util/mixed_dtype_utils.hpp"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
|
||||
#include <cutlass/arch/arch.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
|
||||
#include <cutlass/arch/arch.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
|
||||
+44
-39
@@ -1,11 +1,10 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include "../../torch_utils.h"
|
||||
|
||||
#include "../../dispatch_utils.h"
|
||||
#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
|
||||
#include "quant_conversions.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -105,64 +104,70 @@ __global__ void silu_and_mul_per_block_quant_kernel(
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor& scales, int64_t group_size,
|
||||
std::optional<torch::Tensor> scale_ub,
|
||||
void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor& scales,
|
||||
int64_t group_size,
|
||||
std::optional<torch::stable::Tensor> scale_ub,
|
||||
bool is_scale_transposed) {
|
||||
static c10::ScalarType kFp8Type = is_fp8_ocp()
|
||||
? c10::ScalarType::Float8_e4m3fn
|
||||
: c10::ScalarType::Float8_e4m3fnuz;
|
||||
static torch::headeronly::ScalarType kFp8Type =
|
||||
is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
|
||||
: torch::headeronly::ScalarType::Float8_e4m3fnuz;
|
||||
|
||||
TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
|
||||
TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
|
||||
TORCH_CHECK(
|
||||
input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16,
|
||||
STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
|
||||
out.scalar_type() == torch::headeronly::ScalarType::Char);
|
||||
STD_TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
|
||||
STD_TORCH_CHECK(
|
||||
input.scalar_type() == torch::headeronly::ScalarType::Half ||
|
||||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
|
||||
"Input must be FP16 or BF16");
|
||||
TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
|
||||
TORCH_CHECK(group_size == 128 || group_size == 64,
|
||||
"Unsupported group size: ", group_size);
|
||||
STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
|
||||
STD_TORCH_CHECK(group_size == 128 || group_size == 64,
|
||||
"Unsupported group size: ", group_size);
|
||||
|
||||
if (scale_ub.has_value()) {
|
||||
TORCH_CHECK(out.dtype() == kFp8Type);
|
||||
STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
|
||||
}
|
||||
|
||||
int32_t hidden_size = out.size(-1);
|
||||
auto num_tokens = input.size(0);
|
||||
int32_t num_groups = hidden_size / group_size;
|
||||
|
||||
TORCH_CHECK(input.size(-1) == hidden_size * 2,
|
||||
"input last dim must be 2x output hidden_size");
|
||||
TORCH_CHECK(hidden_size % group_size == 0,
|
||||
"hidden_size must be divisible by group_size");
|
||||
STD_TORCH_CHECK(input.size(-1) == hidden_size * 2,
|
||||
"input last dim must be 2x output hidden_size");
|
||||
STD_TORCH_CHECK(hidden_size % group_size == 0,
|
||||
"hidden_size must be divisible by group_size");
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
dim3 grid(num_tokens, num_groups);
|
||||
dim3 block(group_size);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
|
||||
using scalar_in_t = scalar_t;
|
||||
|
||||
VLLM_DISPATCH_QUANT_TYPES(
|
||||
VLLM_STABLE_DISPATCH_QUANT_TYPES(
|
||||
out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
|
||||
using scalar_out_t = scalar_t;
|
||||
|
||||
VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
|
||||
VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
|
||||
vllm::silu_and_mul_per_block_quant_kernel<
|
||||
scalar_in_t, scalar_out_t, transpose_scale, gs>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out.data_ptr<scalar_out_t>(),
|
||||
scales.data_ptr<float>(),
|
||||
input.data_ptr<scalar_in_t>(),
|
||||
scale_ub.has_value() ? scale_ub->data_ptr<float>()
|
||||
: nullptr,
|
||||
hidden_size);
|
||||
});
|
||||
VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
|
||||
VLLM_STABLE_DISPATCH_BOOL(
|
||||
is_scale_transposed, transpose_scale, [&] {
|
||||
vllm::silu_and_mul_per_block_quant_kernel<
|
||||
scalar_in_t, scalar_out_t, transpose_scale, gs>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out.mutable_data_ptr<scalar_out_t>(),
|
||||
scales.mutable_data_ptr<float>(),
|
||||
input.const_data_ptr<scalar_in_t>(),
|
||||
scale_ub.has_value()
|
||||
? scale_ub->const_data_ptr<float>()
|
||||
: nullptr,
|
||||
hidden_size);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -20,7 +20,7 @@
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
namespace vllm::c3x {
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
/*
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
#include "cuda_utils.h"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
|
||||
void dispatch_scaled_mm(torch::stable::Tensor& c,
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
#include "get_group_starts.cuh"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
@@ -23,7 +23,7 @@
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
using namespace cute;
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
void cutlass_scaled_mm_sm75(torch::stable::Tensor& c,
|
||||
torch::stable::Tensor const& a,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
#include "ops.h"
|
||||
#include "cuda_utils.h"
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
@@ -26,9 +27,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
|
||||
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
|
||||
"()");
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
@@ -321,6 +325,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor? scale_ub, Tensor!? residual, int group_size, "
|
||||
"bool is_scale_transposed) -> ()");
|
||||
|
||||
// Fused SiLU+Mul + per-block quantization
|
||||
ops.def(
|
||||
"silu_and_mul_per_block_quant("
|
||||
"Tensor! out, "
|
||||
"Tensor input, "
|
||||
"Tensor! scales, "
|
||||
"int group_size, "
|
||||
"Tensor? scale_ub=None, "
|
||||
"bool is_scale_transposed=False) -> ()");
|
||||
|
||||
// Rotary embedding
|
||||
// Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
|
||||
ops.def(
|
||||
@@ -343,6 +357,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
|
||||
"int q_head_padded, float eps, int cache_block_size) -> Tensor");
|
||||
|
||||
// FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
|
||||
// FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
|
||||
ops.def(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert("
|
||||
"Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, "
|
||||
"Tensor position_ids, Tensor cos_sin_cache, float eps, "
|
||||
"int cache_block_size) -> ()");
|
||||
ops.def(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert("
|
||||
"Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, "
|
||||
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
|
||||
"Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
|
||||
"int cache_block_size) -> ()");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.def(
|
||||
"minimax_allreduce_rms("
|
||||
@@ -533,9 +561,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("per_token_group_quant_int8",
|
||||
TORCH_BOX(&per_token_group_quant_int8));
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("permute_cols", TORCH_BOX(&permute_cols));
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// CUTLASS scaled_mm ops
|
||||
@@ -585,12 +611,20 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("rms_norm_dynamic_per_token_quant",
|
||||
TORCH_BOX(&rms_norm_dynamic_per_token_quant));
|
||||
ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
|
||||
ops.impl("silu_and_mul_per_block_quant",
|
||||
TORCH_BOX(&silu_and_mul_per_block_quant));
|
||||
|
||||
// Positional encoding kernels (shared CUDA/ROCm)
|
||||
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
|
||||
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
|
||||
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
|
||||
ops.impl(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));
|
||||
ops.impl(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
|
||||
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
|
||||
@@ -641,6 +675,28 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
|
||||
}
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#ifndef USE_ROCM
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
|
||||
ops.impl("get_cuda_view_from_cpu_tensor",
|
||||
TORCH_BOX(&get_cuda_view_from_cpu_tensor));
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
|
||||
cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
|
||||
cuda_utils.def(
|
||||
"get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
|
||||
cuda_utils) {
|
||||
cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
|
||||
cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
|
||||
TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
// These capability-check functions take only primitive args (no tensors), so
|
||||
// there is no device to dispatch on. CompositeExplicitAutograd makes them
|
||||
// available for all backends. This is the stable ABI equivalent of calling
|
||||
|
||||
+3
-8
@@ -40,12 +40,6 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
|
||||
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
|
||||
torch::Tensor& weight, double epsilon);
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor& scales, int64_t group_size,
|
||||
std::optional<torch::Tensor> scale_ub,
|
||||
bool is_scale_transposed);
|
||||
|
||||
// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
|
||||
// ABI for CUDA). It remains here because the CPU build still uses these
|
||||
// torch::Tensor declarations.
|
||||
@@ -84,8 +78,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
|
||||
torch::Tensor const& seq_lens,
|
||||
torch::Tensor const& page_table, double scale);
|
||||
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
|
||||
|
||||
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
torch::Tensor const& scale,
|
||||
std::optional<torch::Tensor> const& azp);
|
||||
@@ -110,4 +102,7 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
|
||||
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
||||
int64_t quant_level, bool cast_bf2half = false);
|
||||
int64_t qr_max_size();
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,639 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100).
|
||||
//
|
||||
// Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3
|
||||
// W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch.
|
||||
// Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering
|
||||
// a tile of N output columns and K input positions.
|
||||
//
|
||||
// Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled,
|
||||
// [E, groups, N] scales, [E, groups, N/8] packed zeros.
|
||||
//
|
||||
// Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread
|
||||
// handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to
|
||||
// the pre-zeroed output tensor (no FP32 scratch buffer).
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hip/hip_bf16.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
|
||||
#include "qdq_4_rdna3.cuh"
|
||||
|
||||
#if defined(__HIPCC__) && defined(__gfx1100__)
|
||||
#define __HIP__RDNA3__
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
namespace moe_gptq_rdna3 {
|
||||
|
||||
#define BLOCK_KN_SIZE 256
|
||||
#define THREADS_X 256
|
||||
|
||||
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
|
||||
|
||||
using gptq_rdna3::bf162_t;
|
||||
using gptq_rdna3::bf16_t;
|
||||
|
||||
// --- Helpers (same as q_gemm_rdna3.cu) ---
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ T tzero();
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ half tzero<half>() {
|
||||
return __float2half_rn(0.0f);
|
||||
}
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ bf16_t tzero<bf16_t>() {
|
||||
return __float2bfloat16(0.0f);
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
|
||||
float result = 0.0f;
|
||||
const half2* a2_ptr = (const half2*)a_ptr;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
|
||||
const bf16_t* a_ptr) {
|
||||
float result = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
uint32_t aw;
|
||||
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
|
||||
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
|
||||
float a_y = __uint_as_float(aw & 0xFFFF0000u);
|
||||
result = __fmaf_rn(dq[2 * i + 0], a_x, result);
|
||||
result = __fmaf_rn(dq[2 * i + 1], a_y, result);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
|
||||
half2 v23) {
|
||||
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
|
||||
unsigned long long old = *addr_u;
|
||||
while (true) {
|
||||
union {
|
||||
unsigned long long u;
|
||||
half2 h2[2];
|
||||
} cur, sum;
|
||||
cur.u = old;
|
||||
sum.h2[0] = __hadd2(cur.h2[0], v01);
|
||||
sum.h2[1] = __hadd2(cur.h2[1], v23);
|
||||
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
|
||||
if (prev == old) break;
|
||||
old = prev;
|
||||
}
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
|
||||
bf162_t v23) {
|
||||
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
|
||||
unsigned long long old = *addr_u;
|
||||
while (true) {
|
||||
union {
|
||||
unsigned long long u;
|
||||
bf162_t b2[2];
|
||||
} cur, sum;
|
||||
cur.u = old;
|
||||
sum.b2[0] = __hadd2(cur.b2[0], v01);
|
||||
sum.b2[1] = __hadd2(cur.b2[1], v23);
|
||||
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
|
||||
if (prev == old) break;
|
||||
old = prev;
|
||||
}
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
|
||||
int (&zeros)[4]) {
|
||||
int qcol = n / 8;
|
||||
int shift = (n & 0x07) * 4;
|
||||
uint32_t d = qzeros_row[qcol] >> shift;
|
||||
zeros[0] = (int)(d & 0xF);
|
||||
zeros[1] = (int)((d >> 4) & 0xF);
|
||||
zeros[2] = (int)((d >> 8) & 0xF);
|
||||
zeros[3] = (int)((d >> 12) & 0xF);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
|
||||
T (&scales)[4]) {
|
||||
scales[0] = scales_row[n + 0];
|
||||
scales[1] = scales_row[n + 1];
|
||||
scales[2] = scales_row[n + 2];
|
||||
scales[3] = scales_row[n + 3];
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Fused MoE kernel.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <typename T, int BLOCK_SIZE_M>
|
||||
__global__ void moe_gemm_q4_kernel_rdna3(
|
||||
const T* __restrict__ a, // [size_m, size_k] or [M*topk, K]
|
||||
T* __restrict__ c, // [M*topk, size_n] pre-zeroed
|
||||
const uint32_t* __restrict__ b_q_weight, // [E, K/8, N] packed
|
||||
const T* __restrict__ b_scales, // [E, groups, N]
|
||||
const uint32_t* __restrict__ b_qzeros, // [E, groups, N/8] packed
|
||||
const float* __restrict__ topk_weights, // [M*topk] or nullptr
|
||||
const int32_t* __restrict__ sorted_token_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ num_tokens_post_padded,
|
||||
const int size_m, // total tokens (original M, or M*topk for w2)
|
||||
const int size_n, // output features per expert
|
||||
const int size_k, // input features
|
||||
const int groups, // K / group_size
|
||||
const int top_k, // routing top-k (1 for w2 pass)
|
||||
// Per-expert strides (in elements, not bytes)
|
||||
const int expert_weight_stride, // (K/8) * N
|
||||
const int expert_scales_stride, // groups * N
|
||||
const int expert_zeros_stride, // groups * (N/8)
|
||||
const bool mul_topk_weight,
|
||||
const int output_topk) { // >0: reduce output by token_id/output_topk
|
||||
const int t = threadIdx.x;
|
||||
const int token_block = blockIdx.x;
|
||||
const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4;
|
||||
const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
|
||||
const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
|
||||
const int n = offset_n + t * 4;
|
||||
|
||||
// Early exit for padding blocks or invalid experts (expert_map = -1)
|
||||
if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return;
|
||||
|
||||
const int expert_id = expert_ids[token_block];
|
||||
if (expert_id == -1) return;
|
||||
|
||||
// Expert-specific pointers
|
||||
const uint32_t* expert_weights =
|
||||
b_q_weight + (int64_t)expert_id * expert_weight_stride;
|
||||
const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride;
|
||||
const uint32_t* expert_qzeros =
|
||||
b_qzeros + (int64_t)expert_id * expert_zeros_stride;
|
||||
|
||||
// LDS for activations
|
||||
constexpr int LDS_PAD = 8;
|
||||
__shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD];
|
||||
|
||||
static_assert(BLOCK_KN_SIZE == THREADS_X,
|
||||
"BLOCK_KN_SIZE must equal THREADS_X");
|
||||
|
||||
// For bf16 M=1, we can skip LDS and read A from global (same as dense).
|
||||
// fp16 always needs LDS due to the dot22_8_f indexing pattern.
|
||||
constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
|
||||
|
||||
const int offset_m_base = token_block * BLOCK_SIZE_M;
|
||||
|
||||
if constexpr (USE_LDS_A) {
|
||||
if (offset_k + t < end_k) {
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
int32_t token_id = sorted_token_ids[offset_m_base + m];
|
||||
int token_row = token_id / top_k;
|
||||
T av;
|
||||
if (token_row < size_m) {
|
||||
av = a[(int64_t)token_row * size_k + offset_k + t];
|
||||
} else {
|
||||
av = tzero<T>();
|
||||
}
|
||||
block_a[m][t] = av;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (n >= size_n) return;
|
||||
|
||||
// Group bookkeeping
|
||||
const int groupsize = size_k / groups;
|
||||
int group = offset_k / groupsize;
|
||||
int nextgroup = (group + 1) * groupsize;
|
||||
|
||||
// Weight pointer for this expert
|
||||
int qk = offset_k / 8;
|
||||
const uint32_t* b_ptr = expert_weights + qk * size_n + n;
|
||||
|
||||
// Per-column dequant constants (4 columns per thread)
|
||||
half2 z1z16_h[4][2], y1y16_h[4][2];
|
||||
float z_b_f[4], y_b_f[4];
|
||||
|
||||
// GPTQv1: zero_offset = 1
|
||||
constexpr int zero_offset = 1;
|
||||
|
||||
auto refresh_group = [&](int g) {
|
||||
const uint32_t* qz_row = expert_qzeros + g * (size_n / 8);
|
||||
const T* sc_row = expert_scales + g * size_n;
|
||||
int zeros[4];
|
||||
T scales[4];
|
||||
load4_zeros(qz_row, n, zeros);
|
||||
load4_scales<T>(sc_row, n, scales);
|
||||
if constexpr (std::is_same<T, half>::value) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset),
|
||||
scales[i], z1z16_h[i], y1y16_h[i]);
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset),
|
||||
scales[i], z_b_f[i], y_b_f[i]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
refresh_group(group);
|
||||
|
||||
float block_c[BLOCK_SIZE_M][4];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
|
||||
}
|
||||
|
||||
// --- Main K-loop ---
|
||||
int k = offset_k;
|
||||
while (k < end_k) {
|
||||
if (k == nextgroup) {
|
||||
group++;
|
||||
nextgroup += groupsize;
|
||||
refresh_group(group);
|
||||
}
|
||||
|
||||
// Prefetch 4 weight words (128 bytes)
|
||||
int4 b_w[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
b_w[j] = *(const int4*)(b_ptr + j * size_n);
|
||||
}
|
||||
b_ptr += 4 * size_n;
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
const int a_off = (k - offset_k) + 8 * j;
|
||||
|
||||
if constexpr (std::is_same<T, half>::value) {
|
||||
// fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16
|
||||
half2 dq[4][4];
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0],
|
||||
y1y16_h[0]);
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1],
|
||||
y1y16_h[1]);
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2],
|
||||
y1y16_h[2]);
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3],
|
||||
y1y16_h[3]);
|
||||
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
|
||||
block_c[m][0] += dot22_8_f(dq[0], a_ptr);
|
||||
block_c[m][1] += dot22_8_f(dq[1], a_ptr);
|
||||
block_c[m][2] += dot22_8_f(dq[2], a_ptr);
|
||||
block_c[m][3] += dot22_8_f(dq[3], a_ptr);
|
||||
}
|
||||
} else if constexpr (BLOCK_SIZE_M == 1) {
|
||||
// bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
|
||||
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
|
||||
constexpr uint32_t BF16_MAGIC = 0x43004300u;
|
||||
constexpr uint32_t BF16_ONES = 0x3F803F80u;
|
||||
union pack4 {
|
||||
float f[4];
|
||||
uint32_t u[4];
|
||||
};
|
||||
|
||||
uint32_t w[4];
|
||||
__builtin_memcpy(w, &b_w[j], sizeof(int4));
|
||||
|
||||
// Load activations — read from global (no LDS for bf16 M=1)
|
||||
pack4 a_pack;
|
||||
{
|
||||
int32_t token_id = sorted_token_ids[offset_m_base];
|
||||
int token_row = token_id / top_k;
|
||||
if (token_row < size_m) {
|
||||
const uint32_t* a_words = reinterpret_cast<const uint32_t*>(
|
||||
a + (int64_t)token_row * size_k + offset_k + a_off);
|
||||
a_pack.u[0] = a_words[0];
|
||||
a_pack.u[1] = a_words[1];
|
||||
a_pack.u[2] = a_words[2];
|
||||
a_pack.u[3] = a_words[3];
|
||||
} else {
|
||||
a_pack.u[0] = 0;
|
||||
a_pack.u[1] = 0;
|
||||
a_pack.u[2] = 0;
|
||||
a_pack.u[3] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// sum_a for bias correction
|
||||
float sum_a = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
sum_a = __builtin_amdgcn_fdot2_f32_bf16(
|
||||
*((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
|
||||
sum_a, /*clamp=*/false);
|
||||
}
|
||||
|
||||
#pragma unroll 1
|
||||
for (int col = 0; col < 4; ++col) {
|
||||
pack4 q_pack;
|
||||
const uint32_t qa = w[col];
|
||||
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
|
||||
|
||||
float partial = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
partial = __builtin_amdgcn_fdot2_f32_bf16(
|
||||
*((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
|
||||
partial, /*clamp=*/false);
|
||||
}
|
||||
|
||||
block_c[0][col] =
|
||||
__fmaf_rn(y_b_f[col], partial,
|
||||
__fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
|
||||
}
|
||||
} else {
|
||||
// bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations
|
||||
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
|
||||
constexpr uint32_t BF16_MAGIC = 0x43004300u;
|
||||
constexpr uint32_t BF16_ONES = 0x3F803F80u;
|
||||
union pack4 {
|
||||
float f[4];
|
||||
uint32_t u[4];
|
||||
};
|
||||
|
||||
uint32_t w[4];
|
||||
__builtin_memcpy(w, &b_w[j], sizeof(int4));
|
||||
|
||||
pack4 a_pack[BLOCK_SIZE_M];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
const uint32_t* a_words =
|
||||
reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
|
||||
a_pack[m].u[0] = a_words[0];
|
||||
a_pack[m].u[1] = a_words[1];
|
||||
a_pack[m].u[2] = a_words[2];
|
||||
a_pack[m].u[3] = a_words[3];
|
||||
}
|
||||
|
||||
float sum_a[BLOCK_SIZE_M];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
float s = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
|
||||
*((const bf16x2_t*)&BF16_ONES),
|
||||
s, /*clamp=*/false);
|
||||
}
|
||||
sum_a[m] = s;
|
||||
}
|
||||
|
||||
#pragma unroll 1
|
||||
for (int col = 0; col < 4; ++col) {
|
||||
pack4 q_pack;
|
||||
const uint32_t qa = w[col];
|
||||
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
|
||||
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
float partial = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
partial = __builtin_amdgcn_fdot2_f32_bf16(
|
||||
*((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
|
||||
partial, /*clamp=*/false);
|
||||
}
|
||||
block_c[m][col] =
|
||||
__fmaf_rn(y_b_f[col], partial,
|
||||
__fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
k += 32;
|
||||
}
|
||||
|
||||
// --- Epilogue: apply topk_weight and atomic-add to output ---
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
int32_t token_id = sorted_token_ids[offset_m_base + m];
|
||||
if (token_id / top_k >= size_m) continue;
|
||||
|
||||
// Apply router weight
|
||||
if (mul_topk_weight && topk_weights != nullptr) {
|
||||
float tw = topk_weights[token_id];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) block_c[m][j] *= tw;
|
||||
}
|
||||
|
||||
// output_topk > 0: reduce by mapping token_id back to original token
|
||||
// (multiple experts write to the same row via atomics)
|
||||
int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk)
|
||||
: (int64_t)token_id;
|
||||
T* out = c + out_row * size_n + n;
|
||||
if constexpr (std::is_same<T, half>::value) {
|
||||
half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
|
||||
__float2half_rn(block_c[m][1]));
|
||||
half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
|
||||
__float2half_rn(block_c[m][3]));
|
||||
atomic_add_pk4_f16(out, r01, r23);
|
||||
} else {
|
||||
bf162_t r01;
|
||||
r01.x = __float2bfloat16(block_c[m][0]);
|
||||
r01.y = __float2bfloat16(block_c[m][1]);
|
||||
bf162_t r23;
|
||||
r23.x = __float2bfloat16(block_c[m][2]);
|
||||
r23.y = __float2bfloat16(block_c[m][3]);
|
||||
atomic_add_pk4_bf16(out, r01, r23);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#else // non-RDNA3: empty stub for symbol parity
|
||||
|
||||
template <typename T, int BLOCK_SIZE_M>
|
||||
__global__ void moe_gemm_q4_kernel_rdna3(
|
||||
const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*,
|
||||
const int32_t*, const int32_t*, const int32_t*, const int, const int,
|
||||
const int, const int, const int, const int, const int, const int,
|
||||
const bool, const int) {}
|
||||
|
||||
#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Launcher
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <typename T, int BLOCK_SIZE_M>
|
||||
void launch_moe_gemm_q4(
|
||||
const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
|
||||
const uint32_t* b_qzeros, const float* topk_weights,
|
||||
const int32_t* sorted_token_ids, const int32_t* expert_ids,
|
||||
const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
|
||||
int size_n, int size_k, int groups, int top_k, int expert_weight_stride,
|
||||
int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight,
|
||||
int output_topk, cudaStream_t stream) {
|
||||
dim3 block(THREADS_X);
|
||||
dim3 grid(num_token_blocks,
|
||||
(size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
|
||||
(size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
|
||||
|
||||
moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k,
|
||||
expert_weight_stride, expert_scales_stride, expert_zeros_stride,
|
||||
mul_topk_weight, output_topk);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void dispatch_moe_gemm_q4(
|
||||
const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
|
||||
const uint32_t* b_qzeros, const float* topk_weights,
|
||||
const int32_t* sorted_token_ids, const int32_t* expert_ids,
|
||||
const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
|
||||
int size_n, int size_k, int groups, int top_k, int block_size_m,
|
||||
int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride,
|
||||
bool mul_topk_weight, int output_topk, cudaStream_t stream) {
|
||||
// Dispatch to template instantiation based on block_size_m
|
||||
switch (block_size_m) {
|
||||
case 1:
|
||||
launch_moe_gemm_q4<T, 1>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
case 2:
|
||||
launch_moe_gemm_q4<T, 2>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
case 4:
|
||||
launch_moe_gemm_q4<T, 4>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
case 8:
|
||||
launch_moe_gemm_q4<T, 8>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false,
|
||||
"moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, "
|
||||
"got ",
|
||||
block_size_m);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace moe_gptq_rdna3
|
||||
} // namespace vllm
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public entry point
|
||||
// ---------------------------------------------------------------------------
|
||||
//
|
||||
// Inputs:
|
||||
// a [M, K] or [M*top_k, K] half or bfloat16
|
||||
// c [M*top_k, N] same dtype (pre-zeroed!)
|
||||
// b_q_weight [E, K/8, N] uint32 (shuffled)
|
||||
// b_scales [E, groups, N] same dtype as a
|
||||
// b_qzeros [E, groups, N/8] uint32 (packed 4-bit)
|
||||
// topk_weights [M*top_k] or empty float32
|
||||
// sorted_token_ids [num_blocks * block_m] int32
|
||||
// expert_ids [num_blocks] int32
|
||||
// num_tokens_post_padded [1] int32
|
||||
// top_k int
|
||||
// block_size_m int (1, 2, 4, or 8)
|
||||
// mul_topk_weight bool
|
||||
|
||||
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
|
||||
torch::Tensor b_q_weight, torch::Tensor b_scales,
|
||||
torch::Tensor b_qzeros, torch::Tensor topk_weights,
|
||||
torch::Tensor sorted_token_ids,
|
||||
torch::Tensor expert_ids,
|
||||
torch::Tensor num_tokens_post_padded, int64_t top_k,
|
||||
int64_t block_size_m, bool mul_topk_weight,
|
||||
int64_t output_topk) {
|
||||
TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
|
||||
TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor");
|
||||
TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
|
||||
TORCH_CHECK(a.dim() == 2, "a must be 2D");
|
||||
TORCH_CHECK(c.dim() == 2, "c must be 2D");
|
||||
TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]");
|
||||
TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]");
|
||||
TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]");
|
||||
TORCH_CHECK(
|
||||
a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
|
||||
"a must be half or bfloat16");
|
||||
TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
|
||||
"b_scales dtype must match a");
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
int size_m = (int)a.size(0);
|
||||
int size_k = (int)a.size(1);
|
||||
int size_n = (int)b_q_weight.size(2);
|
||||
int groups = (int)b_scales.size(1);
|
||||
|
||||
// Per-expert strides
|
||||
int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));
|
||||
int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2));
|
||||
int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2));
|
||||
|
||||
int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m);
|
||||
|
||||
const float* topk_w_ptr =
|
||||
(topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr;
|
||||
|
||||
// Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't
|
||||
// implicitly convert to half/__hip_bfloat16 in device code).
|
||||
using vllm::gptq_rdna3::bf16_t;
|
||||
|
||||
auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) {
|
||||
using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>;
|
||||
vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
|
||||
a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
|
||||
(const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
|
||||
sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
|
||||
num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m,
|
||||
size_n, size_k, groups, (int)top_k, (int)block_size_m,
|
||||
expert_weight_stride, expert_scales_stride, expert_zeros_stride,
|
||||
mul_topk_weight, (int)output_topk, stream);
|
||||
};
|
||||
|
||||
if (a.scalar_type() == torch::kHalf) {
|
||||
dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(),
|
||||
(const half*)b_scales.data_ptr());
|
||||
} else {
|
||||
dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(),
|
||||
(const bf16_t*)b_scales.data_ptr());
|
||||
}
|
||||
}
|
||||
@@ -27,6 +27,15 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
|
||||
torch::Tensor b_scales,
|
||||
torch::Tensor b_g_idx, bool use_v2_format);
|
||||
|
||||
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
|
||||
torch::Tensor b_q_weight, torch::Tensor b_scales,
|
||||
torch::Tensor b_qzeros, torch::Tensor topk_weights,
|
||||
torch::Tensor sorted_token_ids,
|
||||
torch::Tensor expert_ids,
|
||||
torch::Tensor num_tokens_post_padded, int64_t top_k,
|
||||
int64_t block_size_m, bool mul_topk_weight,
|
||||
int64_t output_topk);
|
||||
|
||||
void paged_attention(
|
||||
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
|
||||
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
|
||||
|
||||
@@ -50,6 +50,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
|
||||
"gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
|
||||
"Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
|
||||
rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
|
||||
|
||||
rocm_ops.def(
|
||||
"moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, "
|
||||
"Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, "
|
||||
"Tensor sorted_token_ids, Tensor expert_ids, "
|
||||
"Tensor num_tokens_post_padded, "
|
||||
"int top_k, int block_size_m, bool mul_topk_weight, "
|
||||
"int output_topk) -> ()");
|
||||
rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3);
|
||||
#endif
|
||||
|
||||
// Custom attention op
|
||||
|
||||
+9
-15
@@ -32,30 +32,24 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
|
||||
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
|
||||
|
||||
#ifdef USE_ROCM
|
||||
// TODO: Remove this once we upgrade to torch 2.11.
|
||||
// ROCm still uses torch 2.10,
|
||||
// So we still need to use unstable torch ABI for now.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
|
||||
&get_cuda_view_from_cpu_tensor);
|
||||
#endif
|
||||
|
||||
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
|
||||
// Fused SiLU+Mul + per-block quantization
|
||||
ops.def(
|
||||
"silu_and_mul_per_block_quant("
|
||||
"Tensor! out, "
|
||||
"Tensor input, "
|
||||
"Tensor! scales, "
|
||||
"int group_size, "
|
||||
"Tensor? scale_ub=None, "
|
||||
"bool is_scale_transposed=False) -> ()");
|
||||
ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
|
||||
&silu_and_mul_per_block_quant);
|
||||
|
||||
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
|
||||
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
|
||||
// kernel launch. Registered in _C_stable_libtorch.
|
||||
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
|
||||
// full-cache bf16/fp8 variants).
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
@@ -175,11 +169,10 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
|
||||
|
||||
custom_ar.def("qr_max_size", &qr_max_size);
|
||||
}
|
||||
#endif
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
|
||||
// Cuda utils
|
||||
|
||||
// Gets the specified device attribute.
|
||||
cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
|
||||
cuda_utils.impl("get_device_attribute", &get_device_attribute);
|
||||
@@ -190,5 +183,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
|
||||
cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
|
||||
&get_max_shared_memory_per_block_device_attribute);
|
||||
}
|
||||
#endif
|
||||
|
||||
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|
||||
|
||||
+22
-22
@@ -255,51 +255,49 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
|
||||
# build stage doesn't need the rust toolchain, protoc, or the rust source.
|
||||
# This stage runs in parallel with csrc-build/extensions-build.
|
||||
FROM ${BUILD_BASE_IMAGE} AS rust-build
|
||||
# This stage reuses the Python environment from base and runs in parallel with
|
||||
# csrc-build/extensions-build.
|
||||
FROM base AS rust-build
|
||||
ARG BUILD_OS
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install a basic C toolchain (some rust crates compile C in their build.rs
|
||||
# scripts) and unzip (used to extract the pinned protoc release below).
|
||||
# Install native tools needed only for Rust/protoc builds.
|
||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
dnf install -y --setopt=install_weak_deps=False \
|
||||
ca-certificates curl git gcc gcc-c++ make unzip \
|
||||
make unzip \
|
||||
&& dnf clean all && rm -rf /var/cache/dnf; \
|
||||
else \
|
||||
apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
make unzip \
|
||||
&& rm -rf /var/lib/apt/lists/*; \
|
||||
fi
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Copy only the rust workspace — the binary is the sole artifact we need.
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
# Build the release binary. Cache cargo registry/git and target/, but copy the
|
||||
# binary out of the target/ cache mount so it persists into the image layer
|
||||
# for later COPY --from=rust-build.
|
||||
# Build the release artifacts. Cache cargo registry/git, but not target/,
|
||||
# because stale target metadata can outlive source updates across BuildKit
|
||||
# cache reuse.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
bash build_rust.sh
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
|
||||
#################### CSRC BUILD IMAGE ####################
|
||||
@@ -342,6 +340,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY pyproject.toml setup.py CMakeLists.txt ./
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY cmake cmake/
|
||||
COPY csrc csrc/
|
||||
COPY vllm/envs.py vllm/envs.py
|
||||
@@ -506,9 +505,10 @@ WORKDIR /workspace
|
||||
COPY --from=csrc-build /workspace/dist /precompiled-wheels
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
ARG GIT_REPO_CHECK=0
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
|
||||
+17
-17
@@ -93,35 +93,34 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Copy only the rust workspace — the binary is the sole artifact we need.
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
# Build the release binary. Cache cargo registry/git and target/, but copy the
|
||||
# binary out of the target/ cache mount so it persists into the image layer
|
||||
# for later COPY --from=rust-build.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
# Build the release artifacts. Cache cargo registry/git, but not target/,
|
||||
# because stale target metadata can outlive source updates across BuildKit
|
||||
# cache reuse.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
|
||||
bash build_rust.sh
|
||||
|
||||
######################### BUILD IMAGE #########################
|
||||
FROM base AS vllm-build
|
||||
@@ -154,9 +153,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
|
||||
@@ -102,21 +102,22 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
@@ -125,8 +126,7 @@ ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
bash build_rust.sh
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
|
||||
#################### WHEEL BUILD IMAGE ####################
|
||||
@@ -139,9 +139,10 @@ ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
|
||||
+19
-17
@@ -98,7 +98,6 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \
|
||||
ARG USE_SCCACHE
|
||||
ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
|
||||
ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
|
||||
ENV SCCACHE_ENDPOINT=${USE_SCCACHE:+${SCCACHE_ENDPOINT}}
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
|
||||
ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
|
||||
|
||||
@@ -139,27 +138,25 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
ENV CARGO_NET_RETRY=10
|
||||
ENV RUSTUP_MAX_RETRIES=10
|
||||
|
||||
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
|
||||
cd ${COMMON_WORKDIR}/vllm \
|
||||
&& uv pip install --system -r requirements/build/rust.txt
|
||||
|
||||
# Build the release binary. Cargo's registry/git caches can be written by
|
||||
# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
|
||||
# keeping the cache benefit. Copy the binary out so it persists into the image
|
||||
# layer for later COPY --from=rust-build.
|
||||
# keeping the cache benefit. Do not cache target/, because stale target metadata
|
||||
# can outlive source updates across BuildKit cache reuse.
|
||||
RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
|
||||
--mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \
|
||||
cd ${COMMON_WORKDIR}/vllm \
|
||||
&& VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
|
||||
&& test -x /tmp/vllm-rs
|
||||
&& bash build_rust.sh \
|
||||
&& test -x vllm/vllm-rs
|
||||
|
||||
# -----------------------
|
||||
# vLLM native build stages
|
||||
@@ -179,6 +176,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
|
||||
# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
|
||||
# not invalidate the expensive native build layer.
|
||||
COPY setup.py CMakeLists.txt ./
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY cmake cmake/
|
||||
COPY csrc csrc/
|
||||
COPY vllm/envs.py vllm/envs.py
|
||||
@@ -210,9 +208,10 @@ ENV VLLM_TARGET_DEVICE=rocm
|
||||
|
||||
COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
|
||||
|
||||
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
|
||||
cd vllm \
|
||||
@@ -231,6 +230,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
|
||||
@@ -418,9 +418,10 @@ FROM fetch_vllm AS build_vllm_wheel_release
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
|
||||
|
||||
# Create /install directory for custom wheels
|
||||
RUN mkdir -p /install
|
||||
@@ -514,6 +515,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requir
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
|
||||
|
||||
+104
-90
@@ -6,21 +6,22 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
@@ -29,8 +30,7 @@ ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
bash build_rust.sh
|
||||
|
||||
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
|
||||
|
||||
@@ -89,8 +89,8 @@ RUN mkdir neo && \
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
|
||||
&& 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.3.
|
||||
@@ -99,8 +99,8 @@ RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
|
||||
RUN rm -f /opt/intel/oneapi/ccl/latest && \
|
||||
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
|
||||
rm -f /opt/intel/oneapi/ccl/latest && \
|
||||
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
|
||||
|
||||
SHELL ["bash", "-c"]
|
||||
@@ -119,24 +119,104 @@ ENV UV_LINK_MODE="copy"
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
|
||||
--mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
|
||||
--mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/xpu.txt && \
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
|
||||
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt
|
||||
uv pip install --upgrade pip
|
||||
|
||||
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
######################### UCX + NIXL BUILD STAGE #########################
|
||||
# Build UCX and NIXL in a dedicated stage so compiler/autotools layers are
|
||||
# never included in the final runtime image (mirrors ROCm's build_rixl stage).
|
||||
FROM vllm-base AS ucx-nixl-build
|
||||
|
||||
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
|
||||
ARG NIXL_VERSION=0.7.0
|
||||
|
||||
# Build-time only: compiler, autotools, and verbs dev headers
|
||||
RUN apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
autoconf \
|
||||
automake \
|
||||
libtool \
|
||||
pkg-config \
|
||||
libibverbs-dev \
|
||||
librdmacm-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Build UCX and produce a NIXL wheel so the final image needs no compiler.
|
||||
# patchelf (installed via uv) is used by the NIXL wheel build to rewrite
|
||||
# RPATH entries, making the wheel portable across stages.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
|
||||
bash autogen.sh && \
|
||||
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
|
||||
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
|
||||
uv pip install --upgrade meson pybind11 patchelf && \
|
||||
uv pip install -r requirements.txt && \
|
||||
PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \
|
||||
LD_LIBRARY_PATH=/tmp/ucx_install/lib \
|
||||
python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \
|
||||
find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \
|
||||
rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \
|
||||
rm -rf /tmp/ucx_source /tmp/nixl_source
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
|
||||
ARG NIXL_VERSION=0.7.0
|
||||
|
||||
# Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
|
||||
# No compiler or autotools are installed in this stage.
|
||||
COPY --from=ucx-nixl-build /tmp/ucx_install /tmp/ucx_install
|
||||
COPY --from=ucx-nixl-build /tmp/nixl_wheels /tmp/nixl_wheels
|
||||
|
||||
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
|
||||
|
||||
# Install RDMA runtime libraries (no build tools) and the pre-built NIXL wheel.
|
||||
# Do not uninstall/reinstall large Python packages here to avoid extra layer
|
||||
# churn; final package resolution remains in the later app install step.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
rdma-core \
|
||||
libibverbs1 \
|
||||
librdmacm1 \
|
||||
libibumad3 \
|
||||
libibmad5 \
|
||||
libmlx5-1 \
|
||||
libmlx4-1 \
|
||||
ibverbs-providers \
|
||||
librdmacm1t64 \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& uv pip install --no-deps /tmp/nixl_wheels/nixl-*.whl \
|
||||
&& rm -rf /tmp/nixl_wheels
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
|
||||
--mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
|
||||
--mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
uv pip install -r /workspace/vllm/requirements/xpu.txt && \
|
||||
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.1 && \
|
||||
uv pip uninstall oneccl oneccl-devel && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
# Keep source-dependent layers near the end so frequent code-only changes
|
||||
# don't invalidate heavy dependency and UCX/NIXL layers.
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
ARG GIT_REPO_CHECK=0
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
@@ -147,75 +227,9 @@ ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
uv pip install --no-build-isolation .
|
||||
uv pip install --no-build-isolation --no-deps .
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# install NIXL and UCX from source code
|
||||
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
|
||||
ARG NIXL_VERSION=0.7.0
|
||||
|
||||
RUN apt-get update && apt-get install -y \
|
||||
pciutils \
|
||||
net-tools \
|
||||
iproute2 \
|
||||
hwloc \
|
||||
numactl \
|
||||
wget \
|
||||
curl \
|
||||
git \
|
||||
build-essential \
|
||||
autoconf \
|
||||
automake \
|
||||
libtool \
|
||||
pkg-config \
|
||||
rdma-core \
|
||||
libibverbs-dev \
|
||||
ibverbs-utils \
|
||||
libibverbs1 \
|
||||
librdmacm-dev \
|
||||
librdmacm1 \
|
||||
libibumad-dev \
|
||||
libibumad3 \
|
||||
libibmad-dev \
|
||||
libibmad5 \
|
||||
infiniband-diags \
|
||||
perftest \
|
||||
ibutils \
|
||||
libmlx5-1 \
|
||||
libmlx4-1 \
|
||||
ibverbs-providers \
|
||||
librdmacm1t64
|
||||
|
||||
ENV PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig:${PKG_CONFIG_PATH}
|
||||
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
|
||||
bash autogen.sh && \
|
||||
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
|
||||
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
|
||||
cd /tmp/nixl_source && \
|
||||
uv pip install --upgrade meson pybind11 patchelf && \
|
||||
uv pip install -r requirements.txt && \
|
||||
uv pip install . && \
|
||||
rm -rf /tmp/ucx_source /tmp/nixl_source
|
||||
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall oneccl oneccl-devel
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 388 KiB After Width: | Height: | Size: 373 KiB |
@@ -405,6 +405,47 @@ vllm bench serve \
|
||||
|
||||
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
|
||||
|
||||
#### BFCL (Tool-Calling) Benchmark
|
||||
|
||||
The Berkeley Function Calling Leaderboard (BFCL) dataset measures serving
|
||||
latency and throughput on realistic tool-calling traffic. Each request
|
||||
carries a per-sample `tools` schema and chat history, so the server must
|
||||
expose `/v1/chat/completions` with an auto-tool-choice parser enabled.
|
||||
The benchmark client always uses the `openai-chat` backend.
|
||||
|
||||
Start a tool-parser-enabled server, then run the bench. For example, with
|
||||
`gpt-oss-20b`:
|
||||
|
||||
```bash
|
||||
# Server
|
||||
vllm serve openai/gpt-oss-20b \
|
||||
--enable-auto-tool-choice \
|
||||
--tool-call-parser openai \
|
||||
--reasoning-parser openai_gptoss
|
||||
|
||||
# Client
|
||||
vllm bench serve \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--model openai/gpt-oss-20b \
|
||||
--dataset-name hf \
|
||||
--dataset-path gorilla-llm/Berkeley-Function-Calling-Leaderboard \
|
||||
--bfcl-categories simple,live_simple,multiple \
|
||||
--num-prompts 200
|
||||
```
|
||||
|
||||
`--bfcl-categories` is a comma-separated list of BFCL v3 category names
|
||||
(without the `BFCL_v3_` prefix or `.json` suffix). Defaults to
|
||||
`simple,live_simple,multiple`. Other supported non-multi-turn categories
|
||||
include `parallel`, `live_parallel`, `parallel_multiple`,
|
||||
`live_parallel_multiple`, `irrelevance`, `live_irrelevance`,
|
||||
`live_relevance`, `java`, `javascript`, and `rest`. Multi-turn categories
|
||||
are not yet supported.
|
||||
|
||||
The dataset class normalizes BFCL's loose schema dialect (`dict` →
|
||||
`object`, `float` → `number`, `tuple` → `array`, `any` → `string`) so
|
||||
modern grammar backends accept the translated tool definitions.
|
||||
|
||||
#### Other HuggingFaceDataset Examples
|
||||
|
||||
```bash
|
||||
|
||||
@@ -296,7 +296,7 @@ llm = LLM(model="Qwen/Qwen3-8B")
|
||||
The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't,
|
||||
vLLM raises a clear `ImportError` at tokenizer load. The override applies to
|
||||
any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`,
|
||||
`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Modes that don't use the HF
|
||||
`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Models that don't use the HF
|
||||
fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag.
|
||||
|
||||
Tokenizer-bound workloads — long shared prefixes, bursty short prompts,
|
||||
|
||||
@@ -101,7 +101,7 @@ vLLM's `pre-commit` hooks will now run automatically every time you commit.
|
||||
Some `pre-commit` hooks only run in CI. If you need to, you can run them locally with:
|
||||
|
||||
```bash
|
||||
pre-commit run --hook-stage manual mypy-3.10
|
||||
pre-commit run --hook-stage manual mypy-3.11
|
||||
```
|
||||
|
||||
### Documentation
|
||||
|
||||
@@ -228,3 +228,17 @@ MLA decode backends are selected using the standard
|
||||
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
### DeepSeek V4 Decode Backends
|
||||
|
||||
DeepSeek V4 sparse MLA uses its own decode backends, selected via
|
||||
`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
|
||||
`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
|
||||
pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
|
||||
default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
|
||||
@@ -82,11 +82,13 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
|
||||
| Architecture | Models | CG for Image | CG for Video |
|
||||
| ------------ | ------ | ------------ | ------------ |
|
||||
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
|
||||
| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
|
||||
| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |
|
||||
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
|
||||
|
||||
@@ -60,7 +60,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
|
||||
- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]
|
||||
- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]
|
||||
- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod]
|
||||
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
|
||||
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.UnquantizedFusedMoEMethod]
|
||||
|
||||
## Fused Experts Kernels
|
||||
|
||||
|
||||
@@ -128,7 +128,7 @@ The lease mechanism is controlled through `kv_connector_extra_config` in `--kv-t
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_connector_extra_config": {"kv_lease_duration": 60}
|
||||
}'
|
||||
```
|
||||
|
||||
@@ -1,319 +0,0 @@
|
||||
# P2P NCCL Connector
|
||||
|
||||
An implementation of xPyD with dynamic scaling based on point-to-point communication, partly inspired by Dynamo.
|
||||
|
||||
## Detailed Design
|
||||
|
||||
### Overall Process
|
||||
|
||||
As shown in Figure 1, the overall process of this **PD disaggregation** solution is described through a request flow:
|
||||
|
||||
1. The client sends an HTTP request to the Proxy/Router's `/v1/completions` interface.
|
||||
2. The Proxy/Router selects a **1P1D (1 Prefill instance + 1 Decode instance)** through either through round-robin or random selection, generates a `request_id` (rules to be introduced later), modifies the `max_tokens` in the HTTP request message to **1**, and then forwards the request to the **P instance**.
|
||||
3. Immediately afterward, the Proxy/Router forwards the **original HTTP request** to the **D instance**.
|
||||
4. The **P instance** performs **Prefill** and then **actively sends the generated KV cache** to the D instance (using **PUT_ASYNC** mode). The D instance's `zmq_addr` can be resolved through the `request_id`.
|
||||
5. The **D instance** has a **dedicated thread** for receiving the KV cache (to avoid blocking the main process). The received KV cache is saved into the **GPU memory buffer**, the size of which is determined by the vLLM startup parameter `kv_buffer_size`. When the GPU buffer is full, the KV cache is stored in the **local Tensor memory pool**.
|
||||
6. During the **Decode**, the D instance's main process retrieves the KV cache (transmitted by the P instance) from either the **GPU buffer** or the **memory pool**, thereby **skipping Prefill**.
|
||||
7. After completing **Decode**, the D instance returns the result to the **Proxy/Router**, which then forwards it to the **client**.
|
||||
|
||||

|
||||
|
||||
### Proxy/Router (Demo)
|
||||
|
||||
A simple HTTP service acts as the entry point for client requests and starts a background thread to listen for P/D instances reporting their HTTP IP and PORT, as well as ZMQ IP and PORT. It maintains a dictionary of `http_addr -> zmq_addr`. The `http_addr` is the IP:PORT for the vLLM instance's request, while the `zmq_addr` is the address for KV cache handshake and metadata reception.
|
||||
|
||||
The Proxy/Router is responsible for selecting 1P1D based on the characteristics of the client request, such as the prompt, and generating a corresponding `request_id`, for example:
|
||||
|
||||
```text
|
||||
cmpl-___prefill_addr_10.0.1.2:21001___decode_addr_10.0.1.3:22001_93923d63113b4b338973f24d19d4bf11-0
|
||||
```
|
||||
|
||||
Currently, to quickly verify whether xPyD can work, a round-robin selection of 1P1D is used. In the future, it is planned to use a trie combined with the load status of instances to select appropriate P and D.
|
||||
|
||||
Each P/D instance periodically sends a heartbeat packet to the Proxy/Router (currently every 3 seconds) to register (i.e., report `http_addr -> zmq_addr`) and keep the connection alive. If an instance crashes and fails to send a ping for a certain period of time, the Proxy/Router will remove the timed-out instance (this feature has not yet been developed).
|
||||
|
||||
### KV Cache Transfer Methods
|
||||
|
||||
There are three methods for KVCache transfer: PUT, GET, and PUT_ASYNC. These methods can be specified using the `--kv-transfer-config` and `kv_connector_extra_config` parameters, specifically through the `send_type` field. Both PUT and PUT_ASYNC involve the P instance actively sending KVCache to the D instance. The difference is that PUT is a synchronous transfer method that blocks the main process, while PUT_ASYNC is an asynchronous transfer method. PUT_ASYNC uses a dedicated thread for sending KVCache, which means it does not block the main process. In contrast, the GET method involves the P instance saving the KVCache to the memory buffer after computing the prefill. The D instance then actively retrieves the computed KVCache from the P instance once it has allocated space for the KVCache.
|
||||
|
||||
Experimental results have shown that the performance of these methods, from highest to lowest, is as follows: PUT_ASYNC → GET → PUT.
|
||||
|
||||
### P2P Communication via ZMQ & NCCL
|
||||
|
||||
As long as the address of the counterpart is known, point-to-point KV cache transfer (using NCCL) can be performed, without being constrained by rank and world size. To support dynamic scaling (expansion and contraction) of instances with PD disaggregation. This means that adding or removing P/D instances does not require a full system restart.
|
||||
|
||||
Each P/D instance only needs to create a single `P2pNcclEngine` instance. This instance maintains a ZMQ Server, which runs a dedicated thread to listen on the `zmq_addr` address and receive control flow requests from other instances. These requests include requests to establish an NCCL connection and requests to send KVCache metadata (such as tensor shapes and data types). However, it does not actually transmit the KVCache data itself.
|
||||
|
||||
When a P instance and a D instance transmit KVCache for the first time, they need to establish a ZMQ connection and an NCCL group. For subsequent KVCache transmissions, this ZMQ connection and NCCL group are reused. The NCCL group consists of only two ranks, meaning the world size is equal to 2. This design is intended to support dynamic scaling, which means that adding or removing P/D instances does not require a full system restart. As long as the address of the counterpart is known, point-to-point KVCache transmission can be performed, without being restricted by rank or world size.
|
||||
|
||||
### NCCL Group Topology
|
||||
|
||||
Currently, only symmetric TP (Tensor Parallelism) methods are supported for KVCache transmission. Asymmetric TP and PP (Pipeline Parallelism) methods will be supported in the future. Figure 2 illustrates the 1P2D setup, where each instance has a TP (Tensor Parallelism) degree of 2. There are a total of 7 NCCL groups: three vLLM instances each have one NCCL group with TP=2. Additionally, the 0th GPU card of the P instance establishes an NCCL group with the 0th GPU card of each D instance. Similarly, the 1st GPU card of the P instance establishes an NCCL group with the 1st GPU card of each D instance.
|
||||
|
||||

|
||||
|
||||
Each NCCL group occupies a certain amount of GPU memory buffer for communication, the size of which is primarily influenced by the `NCCL_MAX_NCHANNELS` environment variable. When `NCCL_MAX_NCHANNELS=16`, an NCCL group typically occupies 100MB, while when `NCCL_MAX_NCHANNELS=8`, it usually takes up 52MB. For large-scale xPyD configurations—such as DeepSeek's 96P144D—this implementation is currently not feasible. Moving forward, we are considering using RDMA for point-to-point communication and are also keeping an eye on UCCL.
|
||||
|
||||
### GPU Memory Buffer and Tensor Memory Pool
|
||||
|
||||
The trade-off in the size of the memory buffer is as follows: For P instances, the memory buffer is not required in PUT and PUT_ASYNC modes, but it is necessary in GET mode. For D instances, a memory buffer is needed in all three modes. The memory buffer for D instances should not be too large. Similarly, for P instances in GET mode, the memory buffer should also not be too large. The memory buffer of D instances is used to temporarily store KVCache sent by P instances. If it is too large, it will reduce the KVCache space available for normal inference by D instances, thereby decreasing the inference batch size and ultimately leading to a reduction in output throughput. The size of the memory buffer is configured by the parameter `kv_buffer_size`, measured in bytes, and is typically set to 5%~10% of the memory size.
|
||||
|
||||
If the `--max-num-seqs` parameter for P instances is set to a large value, due to the large batch size, P instances will generate a large amount of KVCache simultaneously. This may exceed the capacity of the memory buffer of D instances, resulting in KVCache loss. Once KVCache is lost, D instances need to recompute Prefill, which is equivalent to performing Prefill twice. Consequently, the time-to-first-token (TTFT) will significantly increase, leading to degraded performance.
|
||||
|
||||
To address the above issues, I have designed and developed a local Tensor memory pool for storing KVCache, inspired by the buddy system used in Linux memory modules. Since the memory is sufficiently large, typically in the TB range on servers, there is no need to consider prefix caching or using block-based designs to reuse memory, thereby saving space. When the memory buffer is insufficient, KVCache can be directly stored in the Tensor memory pool, and D instances can subsequently retrieve KVCache from it. The read and write speed is that of PCIe, with PCIe 4.0 having a speed of approximately 21 GB/s, which is usually faster than the Prefill speed. Otherwise, solutions like Mooncake and lmcache would not be necessary. The Tensor memory pool acts as a flood diversion area, typically unused except during sudden traffic surges. In the worst-case scenario, my solution performs no worse than the normal situation with a Cache store.
|
||||
|
||||
## Install vLLM
|
||||
|
||||
```shell
|
||||
pip install "vllm>=0.9.2"
|
||||
```
|
||||
|
||||
## Run xPyD
|
||||
|
||||
### Instructions
|
||||
|
||||
- The following examples are run on an A800 (80GB) device, using the Meta-Llama-3.1-8B-Instruct model.
|
||||
- Pay attention to the setting of the `kv_buffer_size` (in bytes). The empirical value is 10% of the GPU memory size. This is related to the kvcache size. If it is too small, the GPU memory buffer for temporarily storing the received kvcache will overflow, causing the kvcache to be stored in the tensor memory pool, which increases latency. If it is too large, the kvcache available for inference will be reduced, leading to a smaller batch size and decreased throughput.
|
||||
- For Prefill instances, when using non-GET mode, the `kv_buffer_size` can be set to 1, as Prefill currently does not need to receive kvcache. However, when using GET mode, a larger `kv_buffer_size` is required because it needs to store the kvcache sent to the D instance.
|
||||
- You may need to modify the `kv_buffer_size` and `port` in the following commands (if there is a conflict).
|
||||
- `PUT_ASYNC` offers the best performance and should be prioritized.
|
||||
- The `--port` must be consistent with the `http_port` in the `--kv-transfer-config`.
|
||||
- The `disagg_proxy_p2p_nccl_xpyd.py` script will use port 10001 (for receiving client requests) and port 30001 (for receiving service discovery from P and D instances).
|
||||
- The node running the proxy must have `quart` installed.
|
||||
- Supports multiple nodes; you just need to modify the `proxy_ip` and `proxy_port` in `--kv-transfer-config`.
|
||||
- In the following examples, it is assumed that **the proxy's IP is 10.0.1.1**.
|
||||
|
||||
### Run 1P3D
|
||||
|
||||
#### Proxy (e.g. 10.0.1.1)
|
||||
|
||||
```shell
|
||||
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
```
|
||||
|
||||
#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20001 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode1 (e.g. 10.0.1.3 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode2 (e.g. 10.0.1.4 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20003 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode3 (e.g. 10.0.1.5 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20004 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
### Run 3P1D
|
||||
|
||||
#### Proxy (e.g. 10.0.1.1)
|
||||
|
||||
```shell
|
||||
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
```
|
||||
|
||||
#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20001 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Prefill2 (e.g. 10.0.1.3 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Prefill3 (e.g. 10.0.1.4 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20003 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode1 (e.g. 10.0.1.5 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20004 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
## Single request
|
||||
|
||||
```shell
|
||||
curl -X POST -s http://10.0.1.1:10001/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "base_model",
|
||||
"prompt": "San Francisco is a",
|
||||
"max_tokens": 10,
|
||||
"temperature": 0
|
||||
}'
|
||||
```
|
||||
|
||||
## Benchmark
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model base_model \
|
||||
--tokenizer meta-llama/Llama-3.1-8B-Instruct \
|
||||
--dataset-name "random" \
|
||||
--host 10.0.1.1 \
|
||||
--port 10001 \
|
||||
--random-input-len 1024 \
|
||||
--random-output-len 1024 \
|
||||
--ignore-eos \
|
||||
--burstiness 100 \
|
||||
--percentile-metrics "ttft,tpot,itl,e2el" \
|
||||
--metric-percentiles "90,95,99" \
|
||||
--seed $(date +%s) \
|
||||
--trust-remote-code \
|
||||
--request-rate 3 \
|
||||
--num-prompts 1000
|
||||
```
|
||||
|
||||
## Shut down
|
||||
|
||||
```shell
|
||||
pgrep python | xargs kill -9 && pkill -f python
|
||||
```
|
||||
|
||||
## Test data
|
||||
|
||||
### **Scenario**: 1K input & 200 output tokens, E2E P99 latency ~2s
|
||||
|
||||

|
||||
@@ -88,7 +88,7 @@ If compilation fails for a multimodal model:
|
||||
|
||||
1. **Disable and test**: First verify the model works without compilation:
|
||||
```bash
|
||||
VLLM_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'
|
||||
vllm serve <model> --compilation-config='{"mode":0,"compile_mm_encoder":"false"}'
|
||||
```
|
||||
|
||||
2. **Check logs**: Enable debug logging to see compilation details:
|
||||
|
||||
@@ -17,10 +17,7 @@ Batch invariance is crucial for several use cases:
|
||||
|
||||
## Hardware Requirements
|
||||
|
||||
Batch invariance currently requires NVIDIA GPUs with compute capability 9.0 or higher:
|
||||
|
||||
- **H-series**: H100, H200
|
||||
- **B-series**: B100, B200
|
||||
Batch invariance requires NVIDIA GPUs with compute capability 8.0 or higher.
|
||||
|
||||
## Enabling Batch Invariance
|
||||
|
||||
@@ -107,7 +104,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **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`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8`
|
||||
- **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`
|
||||
- **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
- **Mistral**: `mistralai/Mistral-7B-v0.3`
|
||||
|
||||
|
||||
@@ -17,8 +17,6 @@ Two main reasons:
|
||||
|
||||
## Usage example
|
||||
|
||||
Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
|
||||
|
||||
Now supports 9 types of connectors:
|
||||
|
||||
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
|
||||
@@ -29,7 +27,6 @@ Now supports 9 types of connectors:
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
|
||||
```
|
||||
|
||||
- **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
|
||||
- **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
|
||||
- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
|
||||
- **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
|
||||
@@ -44,16 +41,14 @@ Now supports 9 types of connectors:
|
||||
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
|
||||
```
|
||||
|
||||
For multi-tier offloading (e.g., CPU + filesystem tier) and the full configuration reference, see the [KV Offloading Usage Guide](kv_offloading_usage.md).
|
||||
|
||||
- **FlexKVConnectorV1**: refer to [examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py](../../examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
|
||||
```
|
||||
|
||||
## Benchmarks
|
||||
|
||||
Please refer to [benchmarks/disagg_benchmarks](../../benchmarks/disagg_benchmarks) for disaggregated prefilling benchmarks.
|
||||
|
||||
## Development
|
||||
|
||||
We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
# KV Offloading Usage Guide
|
||||
|
||||
This guide covers configuration of the [`OffloadingConnector`](disagg_prefill.md), which extends the prefix cache by offloading completed KV blocks to slower but larger tiers (CPU host memory, plus optional secondary tiers) as they are produced. Hits in the offload tiers are promoted back to GPU on demand. Transfers between GPU and CPU use DMA (`cudaMemcpyAsync`) and run asynchronously alongside model computation, so offloading adds minimal CPU- and GPU-core overhead.
|
||||
|
||||
!!! note
|
||||
The `OffloadingConnector` currently supports CUDA, ROCm, and XPU only.
|
||||
|
||||
## Overview
|
||||
|
||||
Two specs are available, selected by the `spec_name` key in `kv_connector_extra_config`:
|
||||
|
||||
- `CPUOffloadingSpec` (default): single CPU tier. Completed GPU blocks are copied into pinned host memory.
|
||||
- `TieringOffloadingSpec`: multi-tier. A CPU primary tier plus one or more secondary tiers.
|
||||
|
||||
Only the CPU primary tier has direct GPU access. Secondary tiers cannot read from or write to GPU memory; all GPU↔secondary transfers are staged through the CPU primary tier.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
GPU <--> CPU["CPU primary tier"]
|
||||
CPU <--> S0["Secondary tier 0"]
|
||||
CPU <--> S1["Secondary tier 1"]
|
||||
CPU <--> SN["..."]
|
||||
```
|
||||
|
||||
## Single-Tier Setup (CPU Only)
|
||||
|
||||
```bash
|
||||
vllm serve <model> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "OffloadingConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {
|
||||
"block_size": 64,
|
||||
"cpu_bytes_to_use": 1000000000
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## Multi-Tier Setup
|
||||
|
||||
Set `spec_name` to `"TieringOffloadingSpec"` and supply a `secondary_tiers` list. Each entry is a dict with a required `type` key plus tier-specific fields. The list is ordered: tier 0 is consulted before tier 1, and so on. See [Secondary Tiers](#secondary-tiers) for tier-specific keys.
|
||||
|
||||
```bash
|
||||
vllm serve <model> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "OffloadingConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {
|
||||
"spec_name": "TieringOffloadingSpec",
|
||||
"cpu_bytes_to_use": 10737418240,
|
||||
"block_size": 16,
|
||||
"eviction_policy": "lru",
|
||||
"secondary_tiers": [
|
||||
{
|
||||
"type": "fs",
|
||||
"root_dir": "/mnt/kv_cache",
|
||||
"n_read_threads": 32,
|
||||
"n_write_threads": 16
|
||||
}
|
||||
]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## `kv_connector_extra_config` Reference
|
||||
|
||||
| Key | Required | Default | Scope | Notes |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
|
||||
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
|
||||
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
|
||||
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
|
||||
Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
|
||||
|
||||
### Filesystem (FS)
|
||||
|
||||
The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
|
||||
|
||||
| Key | Required | Default | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
| `type` | yes | — | Must be `fs`. |
|
||||
| `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). |
|
||||
| `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
|
||||
| `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
|
||||
|
||||
Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
|
||||
|
||||
#### On-Disk Layout
|
||||
|
||||
Under `root_dir`, vLLM creates a subdirectory `<model>_<digest>`, where `<model>` is the model name with `/` replaced by `_` (so HuggingFace IDs like `meta-llama/Llama-3-8B` don't nest), and `<digest>` is a short SHA256 prefix derived from the run configuration (model, block size, parallelism, dtype, etc.). Runs with the same configuration share the same subdirectory; runs with different configurations live side-by-side under the same `root_dir` without colliding.
|
||||
|
||||
Inside that subdirectory, blocks are sharded across hash-prefix subdirectories to limit directory fan-out:
|
||||
|
||||
```text
|
||||
<root_dir>/
|
||||
<model>_<digest>/
|
||||
config.json
|
||||
<model>_<digest>_r<rank>/
|
||||
<hhh>/ # first 3 hex chars of the block hash
|
||||
<hh>_g<group_idx>/ # next 2 hex chars + KV cache group index
|
||||
<hash_hex>.bin # full block hash (in hex)
|
||||
```
|
||||
|
||||
`config.json` records the run (block size, number of KV groups, etc.) and is written on first start. Each rank writes blocks under its own `_r<rank>` sibling directory, so multiple ranks can safely share the same `root_dir`.
|
||||
|
||||
#### Cross-Process Sharing
|
||||
|
||||
To enable KV cache sharing between multiple vLLM instances using the same `root_dir` (e.g., via a shared PVC), the `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g., `"0"`) on every instance. Without this, each process initializes `NONE_HASH` (the chain-hash seed for block content hashes) with random bytes, producing different block filenames for identical token content.
|
||||
|
||||
```bash
|
||||
PYTHONHASHSEED=0 vllm serve ...
|
||||
```
|
||||
|
||||
## Tuning Tips
|
||||
|
||||
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
|
||||
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
|
||||
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
|
||||
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
|
||||
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
|
||||
|
||||
## Further Reading
|
||||
|
||||
- [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).
|
||||
@@ -50,7 +50,7 @@ To select a different backend, set `kv_connector_extra_config.backends` in `--kv
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector":"NixlConnector",
|
||||
"kv_role":"kv_both",
|
||||
"kv_role":"kv_producer",
|
||||
"kv_connector_extra_config":{"backends":["LIBFABRIC"]}
|
||||
}'
|
||||
```
|
||||
@@ -60,7 +60,7 @@ You can also pass JSON keys individually using dotted arguments, and you can app
|
||||
```bash
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config.kv_connector NixlConnector \
|
||||
--kv-transfer-config.kv_role kv_both \
|
||||
--kv-transfer-config.kv_role kv_producer \
|
||||
--kv-transfer-config.kv_connector_extra_config.backends+ LIBFABRIC
|
||||
```
|
||||
|
||||
@@ -81,7 +81,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--port 8100 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}'
|
||||
```
|
||||
|
||||
### Consumer (Decoder) Configuration
|
||||
@@ -96,7 +96,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5601 \
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--port 8200 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}'
|
||||
```
|
||||
|
||||
### Proxy Server
|
||||
@@ -212,10 +212,21 @@ sequenceDiagram
|
||||
Enable bidirectional KV transfer by setting `bidirectional_kv_xfer` in `kv_connector_extra_config` on **both** P and D instances:
|
||||
|
||||
```bash
|
||||
# Prefill instance
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_connector_extra_config": {
|
||||
"bidirectional_kv_xfer": true
|
||||
}
|
||||
}'
|
||||
|
||||
# Decode instance
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_connector_extra_config": {
|
||||
"bidirectional_kv_xfer": true
|
||||
}
|
||||
@@ -283,6 +294,21 @@ curl http://localhost:8000/v1/chat/completions \
|
||||
!!! note
|
||||
The `conversation_id` field is a non-standard extension to the OpenAI API. It is consumed by the proxy and not forwarded to the vLLM engine.
|
||||
|
||||
### Benchmarking the multi-turn proxy
|
||||
|
||||
[`benchmarks/multi_turn/benchmark_serving_multi_turn.py`](../../benchmarks/multi_turn/benchmark_serving_multi_turn.py) supports targeting the disaggregated multi-turn proxy with the `--send-conversation-id` flag, which injects a per-conversation `conversation_id` into every request payload so the proxy can key cross-turn KV cache reuse.
|
||||
|
||||
The flag is **off by default** so the benchmark is compatible with strict OpenAI-compatible frontends that reject unknown top-level fields. When benchmarking the multi-turn proxy you must pass it explicitly — otherwise every turn lands as a cache MISS and the bidirectional KV transfer path is never exercised.
|
||||
|
||||
```bash
|
||||
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \
|
||||
--model <MODEL> --served-model-name <NAME> \
|
||||
--url http://<proxy_host>:8000 \
|
||||
--input-file benchmarks/multi_turn/generate_multi_turn.json \
|
||||
--num-clients 2 --max-active-conversations 6 \
|
||||
--send-conversation-id
|
||||
```
|
||||
|
||||
### Limitations
|
||||
|
||||
- Requires a stateful proxy (or equivalent router) to track and forward `kv_transfer_params` between turns.
|
||||
@@ -359,11 +385,10 @@ For multi-host DP deployment, only need to provide the host/port of the head ins
|
||||
|
||||
- **kv_producer**: For prefiller instances that generate KV caches
|
||||
- **kv_consumer**: For decoder instances that consume KV caches from prefiller
|
||||
- **kv_both**: Enables symmetric functionality where the connector can act as both producer and consumer. This provides flexibility for experimental setups and scenarios where the role distinction is not predetermined.
|
||||
- **kv_both** (deprecated): Previously used as a catch-all when the role was not predetermined. This value is now deprecated for NixlConnector and will be removed in a future release.
|
||||
|
||||
!!! tip
|
||||
NixlConnector currently does not distinguish `kv_role`; the actual prefiller/decoder roles are determined by the upper-level proxy (e.g., `toy_proxy_server.py` using `--prefiller-hosts` and `--decoder-hosts`).
|
||||
Therefore, `kv_role` in `--kv-transfer-config` is effectively a placeholder and does not affect NixlConnector's behavior.
|
||||
!!! warning
|
||||
`kv_role="kv_both"` is deprecated for NixlConnector. Please set `kv_role="kv_producer"` for prefill instances and `kv_role="kv_consumer"` for decode instances. See [#33702](https://github.com/vllm-project/vllm/issues/33702) for details.
|
||||
|
||||
### KV Load Failure Policy
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
Quantization trades off model precision for smaller memory footprint, allowing large models to be run on a wider range of devices.
|
||||
|
||||
!!! tip
|
||||
To get started with quantization, see [LLM Compressor](llm_compressor.md), a library for optimizing models for deployment with vLLM that supports FP8, INT8, INT4, and other quantization formats.
|
||||
To get started with quantization, see [LLM Compressor](llm_compressor/README.md), a library for optimizing models for deployment with vLLM that supports FP8, INT8, INT4, and other quantization formats.
|
||||
|
||||
The following are the supported quantization formats for vLLM:
|
||||
|
||||
@@ -12,9 +12,11 @@ The following are the supported quantization formats for vLLM:
|
||||
- [GGUF](gguf.md)
|
||||
- [GPTQModel](gptqmodel.md)
|
||||
- [Intel Neural Compressor](inc.md)
|
||||
- [INT4 W4A16](int4.md)
|
||||
- [INT8 W8A8](int8.md)
|
||||
- [FP8 W8A8](fp8.md)
|
||||
- [LLM Compressor](llm_compressor/README.md)
|
||||
- [FP8 W8A8](llm_compressor/fp8.md)
|
||||
- [INT4 W4A16](llm_compressor/int4.md)
|
||||
- [INT8 W4A8](llm_compressor/int8_w4a8.md)
|
||||
- [INT8 W8A8](llm_compressor/int8_w8a8.md)
|
||||
- [NVIDIA Model Optimizer](modelopt.md)
|
||||
- [Online Quantization](online.md)
|
||||
- [AMD Quark](quark.md)
|
||||
@@ -46,16 +48,17 @@ th:not(:first-child) {
|
||||
}
|
||||
</style>
|
||||
|
||||
| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | x86 CPU |
|
||||
| ------------------------- | ----- | ------ | ------ | --- | ------ | ------- | --------- | ------- |
|
||||
| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ |
|
||||
| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ |
|
||||
| Marlin (GPTQ/AWQ/FP8/FP4) | ❌ | ✅︎* | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ✅︎ |
|
||||
| FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ |
|
||||
| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ |
|
||||
| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | x86 CPU | Arm CPU |
|
||||
| ------------------------- | ----- | ------ | ------ | --- | ------ | ------- | --------- | ------- | ------- |
|
||||
| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ | ❌ |
|
||||
| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ | ❌ |
|
||||
| Marlin (GPTQ/AWQ/FP8/FP4) | ❌ | ✅︎* | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
|
||||
| llm-compressor INT8 (W8A8)| ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ✅︎ | ✅︎ |
|
||||
| llm-compressor INT8 (W4A8)| ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅︎ |
|
||||
| llm-compressor FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
|
||||
| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
|
||||
| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
|
||||
- Volta refers to SM 7.0, Turing to SM 7.5, Ampere to SM 8.0/8.6, Ada to SM 8.9, and Hopper to SM 9.0.
|
||||
- ✅︎ indicates that the quantization method is supported on the specified hardware.
|
||||
|
||||
+25
-25
@@ -21,9 +21,17 @@ The FP8 types typically supported in hardware have two distinct representations,
|
||||
To produce performant FP8 quantized models with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
|
||||
|
||||
```bash
|
||||
pip install llmcompressor
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves three main steps:
|
||||
@@ -57,36 +65,28 @@ For FP8 quantization, we can recover accuracy with simple RTN quantization. We r
|
||||
|
||||
Since simple RTN does not require data for weight quantization and the activations are quantized dynamically, we do not need any calibration data for this quantization flow.
|
||||
|
||||
??? code
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import QuantizationModifier
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import QuantizationModifier
|
||||
# Configure the simple PTQ quantization
|
||||
recipe = QuantizationModifier(
|
||||
targets="Linear",
|
||||
scheme="FP8_DYNAMIC",
|
||||
ignore=["lm_head"],
|
||||
)
|
||||
|
||||
# Configure the simple PTQ quantization
|
||||
recipe = QuantizationModifier(
|
||||
targets="Linear",
|
||||
scheme="FP8_DYNAMIC",
|
||||
ignore=["lm_head"],
|
||||
)
|
||||
# Apply the quantization algorithm.
|
||||
oneshot(model=model, recipe=recipe)
|
||||
|
||||
# Apply the quantization algorithm.
|
||||
oneshot(model=model, recipe=recipe)
|
||||
|
||||
# Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
|
||||
model.save_pretrained(SAVE_DIR)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
# Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
|
||||
model.save_pretrained(SAVE_DIR)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
### 3. Evaluating Accuracy
|
||||
|
||||
Install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Load and run the model in `vllm`:
|
||||
|
||||
```python
|
||||
+64
-68
@@ -12,15 +12,17 @@ Please visit the HF collection of [quantized INT4 checkpoints of popular LLMs re
|
||||
To use INT4 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
|
||||
|
||||
```bash
|
||||
pip install llmcompressor
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.12"
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves four main steps:
|
||||
@@ -52,55 +54,51 @@ When quantizing weights to INT4, you need sample data to estimate the weight upd
|
||||
It's best to use calibration data that closely matches your deployment data.
|
||||
For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
|
||||
### 3. Applying Quantization
|
||||
|
||||
Now, apply the quantization algorithms:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
# Configure the quantization algorithms
|
||||
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
|
||||
|
||||
# Configure the quantization algorithms
|
||||
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A16-G128
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A16-G128
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
This process creates a W4A16 model with weights quantized to 4-bit integers.
|
||||
|
||||
@@ -141,36 +139,34 @@ lm_eval --model vllm \
|
||||
|
||||
The following is an example of an expanded quantization recipe you can tune to your own use case:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationScheme,
|
||||
QuantizationStrategy,
|
||||
QuantizationType,
|
||||
)
|
||||
recipe = GPTQModifier(
|
||||
targets="Linear",
|
||||
config_groups={
|
||||
"config_group": QuantizationScheme(
|
||||
targets=["Linear"],
|
||||
weights=QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
group_size=128,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
actorder="weight",
|
||||
),
|
||||
```python
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationScheme,
|
||||
QuantizationStrategy,
|
||||
QuantizationType,
|
||||
)
|
||||
recipe = GPTQModifier(
|
||||
targets="Linear",
|
||||
config_groups={
|
||||
"config_group": QuantizationScheme(
|
||||
targets=["Linear"],
|
||||
weights=QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
group_size=128,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
actorder="weight",
|
||||
),
|
||||
},
|
||||
ignore=["lm_head"],
|
||||
update_size=NUM_CALIBRATION_SAMPLES,
|
||||
dampening_frac=0.01,
|
||||
)
|
||||
```
|
||||
),
|
||||
},
|
||||
ignore=["lm_head"],
|
||||
update_size=NUM_CALIBRATION_SAMPLES,
|
||||
dampening_frac=0.01,
|
||||
)
|
||||
```
|
||||
|
||||
## Troubleshooting and Support
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
# INT8 W4A8
|
||||
|
||||
vLLM supports quantizing weights to INT4 and activations to INT8 for memory savings and inference acceleration.
|
||||
This quantization method is particularly useful for reducing model size while maintaining good performance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
To use INT8 W4A8 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library.
|
||||
|
||||
```bash
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves four main steps:
|
||||
|
||||
1. Loading the model
|
||||
2. Preparing calibration data
|
||||
3. Applying quantization
|
||||
4. Evaluating accuracy in vLLM
|
||||
|
||||
### 1. Loading the Model
|
||||
|
||||
Load your model and tokenizer using the standard `transformers` AutoModel classes:
|
||||
|
||||
```python
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
|
||||
MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_ID,
|
||||
dtype="auto",
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||||
```
|
||||
|
||||
### 2. Preparing Calibration Data
|
||||
|
||||
When quantizing activations to INT8 and weights to INT4, you need sample data to estimate the activation scales.
|
||||
It's best to use calibration data that closely matches your deployment data.
|
||||
For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
|
||||
### 3. Applying Quantization
|
||||
|
||||
Now, apply the quantization algorithms.
|
||||
|
||||
The following recipes create W4A8 models (int4 weights, int8 activations). On Arm® CPUs, this is accelerated through [KleidiAI](https://github.com/ARM-software/kleidiai).
|
||||
|
||||
Use groupwise for best accuracy, and channelwise for best inference performance.
|
||||
|
||||
=== "Groupwise"
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
|
||||
# Configure the quantization algorithms
|
||||
recipe = [
|
||||
GPTQModifier(
|
||||
targets="Linear",
|
||||
scheme="W4A8",
|
||||
ignore=["lm_head"],
|
||||
dampening_frac=0.01
|
||||
),
|
||||
]
|
||||
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A8-G128-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
=== "Channelwise"
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from compressed_tensors.quantization import QuantizationStrategy, QuantizationType
|
||||
|
||||
scheme = {
|
||||
"targets": ["Linear"],
|
||||
"weights": {
|
||||
"num_bits": 4,
|
||||
"type": QuantizationType.INT,
|
||||
"strategy": QuantizationStrategy.CHANNEL,
|
||||
"symmetric": True,
|
||||
"dynamic": False,
|
||||
"group_size": None,
|
||||
},
|
||||
"input_activations": {
|
||||
"num_bits": 8,
|
||||
"type": QuantizationType.INT,
|
||||
"strategy": QuantizationStrategy.TOKEN,
|
||||
"dynamic": True,
|
||||
"symmetric": False,
|
||||
"observer": None,
|
||||
},
|
||||
"output_activations": None,
|
||||
}
|
||||
|
||||
recipe = [
|
||||
GPTQModifier(
|
||||
targets="Linear",
|
||||
config_groups={"group_0": scheme},
|
||||
ignore=["lm_head"],
|
||||
dampening_frac=0.01,
|
||||
),
|
||||
]
|
||||
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A8-Channelwise-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
### 4. Evaluating Accuracy
|
||||
|
||||
=== "Groupwise"
|
||||
|
||||
After quantization, you can load and run the model in vLLM:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM("./Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token")
|
||||
```
|
||||
|
||||
To evaluate accuracy, you can use `lm_eval`:
|
||||
|
||||
```bash
|
||||
lm_eval --model vllm \
|
||||
--model_args pretrained="./Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token",add_bos_token=true \
|
||||
--tasks gsm8k \
|
||||
--num_fewshot 5 \
|
||||
--limit 250 \
|
||||
--batch_size 'auto'
|
||||
```
|
||||
|
||||
=== "Channelwise"
|
||||
|
||||
After quantization, you can load and run the model in vLLM:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM("./Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token")
|
||||
```
|
||||
|
||||
To evaluate accuracy, you can use `lm_eval`:
|
||||
|
||||
```bash
|
||||
lm_eval --model vllm \
|
||||
--model_args pretrained="./Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token",add_bos_token=true \
|
||||
--tasks gsm8k \
|
||||
--num_fewshot 5 \
|
||||
--limit 250 \
|
||||
--batch_size 'auto'
|
||||
```
|
||||
|
||||
!!! note
|
||||
Quantized models can be sensitive to the presence of the `bos` token. Make sure to include the `add_bos_token=True` argument when running evaluations.
|
||||
|
||||
## Best Practices
|
||||
|
||||
- Start with 512 samples for calibration data (increase if accuracy drops)
|
||||
- Use a sequence length of 2048 as a starting point
|
||||
- Employ the chat template or instruction template that the model was trained with
|
||||
- If you've fine-tuned a model, consider using a sample of your training data for calibration
|
||||
|
||||
## Troubleshooting and Support
|
||||
|
||||
If you encounter any issues or have feature requests, please open an issue on the [vllm-project/llm-compressor](https://github.com/vllm-project/llm-compressor/issues) GitHub repository.
|
||||
+40
-42
@@ -17,15 +17,17 @@ Please visit the HF collection of [quantized INT8 checkpoints of popular LLMs re
|
||||
To use INT8 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
|
||||
|
||||
```bash
|
||||
pip install llmcompressor
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.12"
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves four main steps:
|
||||
@@ -57,26 +59,24 @@ When quantizing activations to INT8, you need sample data to estimate the activa
|
||||
It's best to use calibration data that closely matches your deployment data.
|
||||
For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
@@ -84,33 +84,31 @@ For a general-purpose instruction-tuned model, you can use a dataset like `ultra
|
||||
|
||||
Now, apply the quantization algorithms:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
# Configure the quantization algorithms
|
||||
recipe = [
|
||||
SmoothQuantModifier(smoothing_strength=0.8),
|
||||
GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
|
||||
]
|
||||
|
||||
# Configure the quantization algorithms
|
||||
recipe = [
|
||||
SmoothQuantModifier(smoothing_strength=0.8),
|
||||
GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
|
||||
]
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W8A8-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W8A8-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
This process creates a W8A8 model with weights and activations quantized to 8-bit integers.
|
||||
|
||||
@@ -169,7 +169,7 @@ speculative decoding, breaking down the guarantees into three key areas:
|
||||
> distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
|
||||
> - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
|
||||
> without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../../tests/v1/spec_decode).
|
||||
> verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
|
||||
|
||||
3. **vLLM Logprob Stability**
|
||||
|
||||
@@ -24,10 +24,11 @@ vllm serve google/gemma-4-E2B-it \
|
||||
--speculative-config '{"method":"mtp","model":"gg-hf-am/gemma-4-E2B-it-assistant","num_speculative_tokens":1}'
|
||||
```
|
||||
|
||||
The E2B, E4B, 26B-A4B, and 31B Gemma 4 IT assistant checkpoints are supported
|
||||
when their configuration uses `model_type: gemma4_assistant`. vLLM maps those
|
||||
checkpoints to `Gemma4MTPModel` internally and wires the assistant layers to
|
||||
share KV cache with the target model.
|
||||
The E2B, E4B, 12B, 26B-A4B, and 31B Gemma 4 IT assistant checkpoints are supported.
|
||||
Tower-based variants use `model_type: gemma4_assistant` and the encoder-free
|
||||
Gemma 4 Unified variant (12B) uses `model_type: gemma4_unified_assistant`.
|
||||
vLLM maps both to `Gemma4MTPModel` internally and wires the assistant layers
|
||||
to share KV cache with the target model.
|
||||
|
||||
If an older vLLM release logs `SpeculativeConfig(method='draft_model', ...)`
|
||||
for a Gemma 4 assistant checkpoint, that release is treating the assistant as a
|
||||
|
||||
@@ -504,6 +504,13 @@ Flags: `--tool-call-parser pythonic --chat-template {see_above}`
|
||||
!!! warning
|
||||
Llama's smaller models frequently fail to emit tool calls in the correct format. Results may vary depending on the model.
|
||||
|
||||
## Benchmarking Tool-Calling Performance
|
||||
|
||||
To measure serving latency and throughput on realistic tool-calling traffic,
|
||||
use the BFCL (Berkeley Function Calling Leaderboard) dataset with
|
||||
`vllm bench serve`. See the [BFCL benchmark example](../benchmarking/cli.md#bfcl-tool-calling-benchmark)
|
||||
for the full server + client commands.
|
||||
|
||||
## How to Write a Tool Parser Plugin
|
||||
|
||||
A tool parser plugin is a Python file containing one or more ToolParser implementations. You can write a ToolParser similar to the `Hermes2ProToolParser` in [vllm/tool_parsers/hermes_tool_parser.py](../../vllm/tool_parsers/hermes_tool_parser.py).
|
||||
|
||||
@@ -41,12 +41,12 @@ pip install -v -r requirements/xpu.txt
|
||||
|
||||
```bash
|
||||
pip uninstall -y triton triton-xpu
|
||||
pip install triton-xpu==3.7.0 --extra-index-url https://download.pytorch.org/whl/xpu
|
||||
pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu
|
||||
```
|
||||
|
||||
!!! 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.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).
|
||||
- For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. 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:
|
||||
|
||||
|
||||
@@ -378,7 +378,7 @@ th {
|
||||
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
|
||||
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
|
||||
| `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
|
||||
| `Cohere2MoeForCausalLM` | Command-A+ | `CohereLabs/command-a-plus-05-2026`, etc. | ✅︎ | ✅︎ |
|
||||
| `Cohere2MoeForCausalLM` | North-Mini-Code | `CohereLabs/North-Mini-Code`, etc. | ✅︎ | ✅︎ |
|
||||
| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
|
||||
| `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
|
||||
| `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -550,7 +550,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
|
||||
| `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision, Command-A+ | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, `CohereLabs/command-a-plus-05-2026`, etc. | | ✅︎ |
|
||||
| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano` | | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
|
||||
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -562,12 +562,14 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `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. | | ✅︎ |
|
||||
| `Gemma4UnifiedForConditionalGeneration` | Gemma 4 Unified | T + I<sup>+</sup> + V + A | `google/gemma-4-12B-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` | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | T + A | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
|
||||
| `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` | | |
|
||||
| `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -664,10 +666,16 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
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.
|
||||
- Gemma 4 assistant checkpoints for speculative decoding use vLLM's Gemma
|
||||
- Gemma 4 assistant checkpoints for speculative decoding use vLLM’s Gemma
|
||||
4 MTP path, not generic draft-model speculative decoding. See the
|
||||
[Gemma 4 assistant model MTP example](../features/speculative_decoding/mtp.md#gemma-4-assistant-models).
|
||||
|
||||
!!! note
|
||||
For `Gemma4UnifiedForConditionalGeneration`:
|
||||
- This is the encoder-free Gemma 4 variant (e.g. `gemma-4-12B-it`). Unlike the tower-based `Gemma4ForConditionalGeneration`, it has **no SigLIP vision encoder** and **no audio encoder**. Raw pixel patches are projected directly into LM space via a Dense+LayerNorm pipeline with factorized positional embeddings, and raw audio waveform frames are projected directly through a multimodal embedder.
|
||||
- All modalities (image, video, audio) are supported.
|
||||
- Gemma 4 Unified assistant checkpoints (`model_type: gemma4_unified_assistant`) use the same MTP path as the tower-based variant. See the [Gemma 4 assistant model MTP example](../features/speculative_decoding/mtp.md#gemma-4-assistant-models).
|
||||
|
||||
!!! 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.
|
||||
|
||||
@@ -702,6 +710,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
|
||||
| `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
|
||||
| `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
|
||||
| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ |
|
||||
| `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -24,8 +24,14 @@ echo "Checking pre-commit/pre-run-check status..."
|
||||
MAX_WAIT=300
|
||||
INTERVAL=60
|
||||
ELAPSED=0
|
||||
# Use a GitHub token if provided to raise the API rate limit (60 -> 5000
|
||||
# requests/hour). Set GITHUB_TOKEN in the Read the Docs environment variables.
|
||||
CURL_AUTH=()
|
||||
if [ -n "$GITHUB_TOKEN" ]; then
|
||||
CURL_AUTH=(-H "Authorization: Bearer $GITHUB_TOKEN")
|
||||
fi
|
||||
while :; do
|
||||
RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
|
||||
RAW=$(curl -sS "${CURL_AUTH[@]}" -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
|
||||
HTTP_CODE=$(printf %s "$RAW" | tail -n1)
|
||||
BODY=$(printf %s "$RAW" | sed '$d')
|
||||
if [ "$HTTP_CODE" != "200" ]; then
|
||||
|
||||
@@ -1,127 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This file demonstrates the example usage of disaggregated prefilling
|
||||
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
|
||||
and then transfer the KV cache between them.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
from multiprocessing import Event, Process
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import KVTransferConfig
|
||||
|
||||
|
||||
def run_prefill(prefill_done):
|
||||
# We use GPU 0 for prefill node.
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
|
||||
# The prefill node receives two requests, while the decode node receives
|
||||
# three requests. So the decode node will only receive the KV Cache for
|
||||
# requests 1 and 3. The decode node will use the KV Cache of requests 1
|
||||
# and 3 and do prefilling on request 2.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"Hi, your name is",
|
||||
# The decode node will actually "prefill" this request.
|
||||
"Tell me a very long story",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
|
||||
|
||||
# Using P2pNcclConnector to transmit KV caches between vLLM instances.
|
||||
# This instance is the prefill node (kv_producer, rank 0).
|
||||
# The number of parallel instances for KV cache transfer is set to 2,
|
||||
# as required for P2pNcclConnector.
|
||||
ktc = KVTransferConfig(
|
||||
kv_connector="P2pNcclConnector",
|
||||
kv_role="kv_producer",
|
||||
kv_rank=0,
|
||||
kv_parallel_size=2,
|
||||
)
|
||||
|
||||
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
|
||||
# memory. You may need to adjust the value to fit your GPU.
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
kv_transfer_config=ktc,
|
||||
max_model_len=2000,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
llm.generate(prompts, sampling_params)
|
||||
print("Prefill node is finished.")
|
||||
prefill_done.set()
|
||||
|
||||
# To keep the prefill node running in case the decode node is not done;
|
||||
# otherwise, the script might exit prematurely, causing incomplete decoding.
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
print("Script stopped by user.")
|
||||
|
||||
|
||||
def run_decode(prefill_done):
|
||||
# We use GPU 1 for decode node.
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"Hi, your name is",
|
||||
"Tell me a very long story",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95)
|
||||
|
||||
# Using P2pNcclConnector to transmit KV caches between vLLM instances.
|
||||
# This instance is the decode node (kv_consumer, rank 1).
|
||||
# The number of parallel instances for KV cache transfer is set to 2,
|
||||
# as required for P2pNcclConnector.
|
||||
ktc = KVTransferConfig(
|
||||
kv_connector="P2pNcclConnector",
|
||||
kv_role="kv_consumer",
|
||||
kv_rank=1,
|
||||
kv_parallel_size=2,
|
||||
)
|
||||
|
||||
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
|
||||
# memory. You may need to adjust the value to fit your GPU.
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
kv_transfer_config=ktc,
|
||||
max_model_len=2000,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
# Wait for the producer to start the pipe
|
||||
print("Waiting for prefill node to finish...")
|
||||
prefill_done.wait()
|
||||
|
||||
# At this point when the prefill_done is set, the kv-cache should have been
|
||||
# transferred to this decode node, so we can start decoding.
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
|
||||
def main():
|
||||
prefill_done = Event()
|
||||
prefill_process = Process(target=run_prefill, args=(prefill_done,))
|
||||
decode_process = Process(target=run_decode, args=(prefill_done,))
|
||||
|
||||
# Start prefill node
|
||||
prefill_process.start()
|
||||
|
||||
# Start decode node
|
||||
decode_process.start()
|
||||
|
||||
# Terminate the prefill node when decode is finished
|
||||
decode_process.join()
|
||||
prefill_process.terminate()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,125 +0,0 @@
|
||||
#!/bin/bash
|
||||
# This file demonstrates the example usage of disaggregated prefilling
|
||||
# We will launch 2 vllm instances (1 for prefill and 1 for decode),
|
||||
# and then transfer the KV cache between them.
|
||||
|
||||
set -xe
|
||||
|
||||
echo "🚧🚧 Warning: The usage of disaggregated prefill is experimental and subject to change 🚧🚧"
|
||||
sleep 1
|
||||
|
||||
# meta-llama/Meta-Llama-3.1-8B-Instruct or deepseek-ai/DeepSeek-V2-Lite
|
||||
MODEL_NAME=${HF_MODEL_NAME:-meta-llama/Meta-Llama-3.1-8B-Instruct}
|
||||
|
||||
# Trap the SIGINT signal (triggered by Ctrl+C)
|
||||
trap 'cleanup' INT
|
||||
|
||||
# Cleanup function
|
||||
cleanup() {
|
||||
echo "Caught Ctrl+C, cleaning up..."
|
||||
# Cleanup commands
|
||||
pgrep python | xargs kill -9
|
||||
pkill -f python
|
||||
echo "Cleanup complete. Exiting."
|
||||
exit 0
|
||||
}
|
||||
|
||||
|
||||
if [[ -z "${VLLM_HOST_IP:-}" ]]; then
|
||||
export VLLM_HOST_IP=127.0.0.1
|
||||
echo "Using default VLLM_HOST_IP=127.0.0.1 (override by exporting VLLM_HOST_IP before running this script)"
|
||||
else
|
||||
echo "Using provided VLLM_HOST_IP=${VLLM_HOST_IP}"
|
||||
fi
|
||||
|
||||
|
||||
# install quart first -- required for disagg prefill proxy serve
|
||||
if python3 -c "import quart" &> /dev/null; then
|
||||
echo "Quart is already installed."
|
||||
else
|
||||
echo "Quart is not installed. Installing..."
|
||||
python3 -m pip install quart
|
||||
fi
|
||||
|
||||
# a function that waits vLLM server to start
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -i localhost:${port}/v1/models > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
|
||||
# You can also adjust --kv-ip and --kv-port for distributed inference.
|
||||
|
||||
# prefilling instance, which is the KV producer
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8100 \
|
||||
--max-model-len 100 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--trust-remote-code \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
|
||||
|
||||
# decoding instance, which is the KV consumer
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8200 \
|
||||
--max-model-len 100 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--trust-remote-code \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":"1e10","kv_port":"14580","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8200","send_type":"PUT_ASYNC"}}' &
|
||||
|
||||
# wait until prefill and decode instances are ready
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
|
||||
# launch a proxy server that opens the service at port 8000
|
||||
# the workflow of this proxy:
|
||||
# - send the request to prefill vLLM instance (port 8100), change max_tokens
|
||||
# to 1
|
||||
# - after the prefill vLLM finishes prefill, send the request to decode vLLM
|
||||
# instance
|
||||
# NOTE: the usage of this API is subject to change --- in the future we will
|
||||
# introduce "vllm connect" to connect between prefill and decode instances
|
||||
python3 ../../benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py &
|
||||
sleep 1
|
||||
|
||||
# serve two example requests
|
||||
output1=$(curl -X POST -s http://localhost:8000/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'"$MODEL_NAME"'",
|
||||
"prompt": "San Francisco is a",
|
||||
"max_tokens": 10,
|
||||
"temperature": 0
|
||||
}')
|
||||
|
||||
output2=$(curl -X POST -s http://localhost:8000/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'"$MODEL_NAME"'",
|
||||
"prompt": "Santa Clara is a",
|
||||
"max_tokens": 10,
|
||||
"temperature": 0
|
||||
}')
|
||||
|
||||
|
||||
# Cleanup commands
|
||||
pgrep python | xargs kill -9
|
||||
pkill -f python
|
||||
|
||||
echo ""
|
||||
|
||||
sleep 1
|
||||
|
||||
# Print the outputs of the curl requests
|
||||
echo ""
|
||||
echo "Output of first request: $output1"
|
||||
echo "Output of second request: $output2"
|
||||
|
||||
echo "🎉🎉 Successfully finished 2 test requests! 🎉🎉"
|
||||
echo ""
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user