forked from Karylab-cklius/vllm
Compare commits
154
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
53f6dd5c6f | ||
|
|
99115fcdcd | ||
|
|
1053e248f0 | ||
|
|
b5bcb3ce88 | ||
|
|
b2f9e4caa4 | ||
|
|
831d3848f1 | ||
|
|
fd10e8946d | ||
|
|
ed13deb376 | ||
|
|
99de48e98f | ||
|
|
bf2b45b5d6 | ||
|
|
8112b6c997 | ||
|
|
15d65f8669 | ||
|
|
e3c2fc3b3c | ||
|
|
2b465b2c42 | ||
|
|
3f47a8384d | ||
|
|
04502deca2 | ||
|
|
d2ca3002d9 | ||
|
|
27d7061ef6 | ||
|
|
ef9975d021 | ||
|
|
56c96b0d91 | ||
|
|
59a6b0411d | ||
|
|
dbccc5ae32 | ||
|
|
a89015c6df | ||
|
|
96fa3f42c9 | ||
|
|
81962bb699 | ||
|
|
92e8518d37 | ||
|
|
77cba0259f | ||
|
|
30fbd05537 | ||
|
|
0906123953 | ||
|
|
bc3629b1c4 | ||
|
|
312ea82e75 | ||
|
|
394beb633b | ||
|
|
7f599d7854 | ||
|
|
eb290ab673 | ||
|
|
cbc3a87200 | ||
|
|
afc94523c9 | ||
|
|
8061dc26bd | ||
|
|
fd9d2ede6f | ||
|
|
e09900436c | ||
|
|
5d07e268b1 | ||
|
|
d742856610 | ||
|
|
544cb724c8 | ||
|
|
c314af1abf | ||
|
|
f19ee27e39 | ||
|
|
5f89a03dcb | ||
|
|
53397fbfac | ||
|
|
49f31d7cee | ||
|
|
74d3b799e1 | ||
|
|
29fdeab254 | ||
|
|
ff6173997d | ||
|
|
8de50e46d4 | ||
|
|
da99ffcc13 | ||
|
|
8040ef2426 | ||
|
|
bf4f633b4c | ||
|
|
854c33f380 | ||
|
|
ac87549cbd | ||
|
|
439f336212 | ||
|
|
ffc4f08c8e | ||
|
|
50aa830482 | ||
|
|
f0553889c0 | ||
|
|
fdaa0d9e59 | ||
|
|
0934b26790 | ||
|
|
9e50e1037e | ||
|
|
b5b61c622c | ||
|
|
b68d7ef262 | ||
|
|
7154856f3d | ||
|
|
3f1d40960f | ||
|
|
0da6e7f3d6 | ||
|
|
5559679229 | ||
|
|
da3a252fd1 | ||
|
|
21fd9e85a0 | ||
|
|
8d28b48d01 | ||
|
|
0164022c90 | ||
|
|
30b0714031 | ||
|
|
2e860de498 | ||
|
|
7eca0e1a64 | ||
|
|
7a29a3c54c | ||
|
|
1240c74c0a | ||
|
|
48ebd6f2f1 | ||
|
|
b153ae6089 | ||
|
|
7a6a5b3667 | ||
|
|
0111002323 | ||
|
|
d30b1ecd1b | ||
|
|
dbd80cc031 | ||
|
|
70009fb934 | ||
|
|
ee1d996367 | ||
|
|
6b0103d1c9 | ||
|
|
9321aff536 | ||
|
|
26d725c334 | ||
|
|
7fe6d3c76b | ||
|
|
2e0da24150 | ||
|
+1 |
0b0bd2b5f6 | ||
|
|
33ef67e9fb | ||
|
|
3e74c60b9c | ||
|
|
d1a8ba63d9 | ||
|
|
1423569ff5 | ||
|
|
9a50464698 | ||
|
|
b9b6306ebe | ||
|
|
ca0defa343 | ||
|
|
0b1a8bb1f6 | ||
|
|
fe5145765f | ||
|
|
dbcc1cdd0a | ||
|
|
a82f1b388f | ||
|
|
190be7dad2 | ||
|
|
94682b79f4 | ||
|
|
0ba2aa35a8 | ||
|
|
aaaeda98dc | ||
|
|
d9cd774198 | ||
|
|
70052fb924 | ||
|
|
318b527cc2 | ||
|
|
6a1acac3fe | ||
|
|
213f681f81 | ||
|
|
33c4f3551c | ||
|
|
caa9cad31e | ||
|
|
7513d071bd | ||
|
|
89f6aa3a9e | ||
|
|
84d26b9ee3 | ||
|
|
9e6746b3c7 | ||
|
|
972848f276 | ||
|
|
2279575cd9 | ||
|
|
5d8e90a966 | ||
|
|
e222c33f2f | ||
|
|
c064fa52b6 | ||
|
|
7e51939e25 | ||
|
|
9863102ed9 | ||
|
|
8c13ee5735 | ||
|
|
41798069f3 | ||
|
|
866fea2b99 | ||
|
|
d02df748bf | ||
|
|
453f01783d | ||
|
|
7b40fb9645 | ||
|
|
8eac21a602 | ||
|
|
a454a1dd25 | ||
|
|
833483f357 | ||
|
|
163ecba377 | ||
|
|
589a5b884b | ||
|
|
5c5434e2d8 | ||
|
|
dd72658e7d | ||
|
|
0d77325b10 | ||
|
|
2ac125123a | ||
|
|
7bdf8cc37c | ||
|
|
bf27e34ebb | ||
|
|
275556c35c | ||
|
|
d65acd83d8 | ||
|
|
80c9d5d5e0 | ||
|
|
da54a5bf05 | ||
|
|
1479bd9e9d | ||
|
|
0231dd5467 | ||
|
|
2659467497 | ||
|
|
a49d37c6b9 | ||
|
|
4501a6d56b | ||
|
|
e18f0037a5 | ||
|
|
b354734d17 | ||
|
|
b91a40e729 |
@@ -0,0 +1,26 @@
|
||||
group: Benchmarks
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 40
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
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/benchmarks/
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s benchmarks/'
|
||||
@@ -2,6 +2,44 @@ group: Engine Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: Engine
|
||||
key: engine
|
||||
timeout_in_minutes: 40
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/config/
|
||||
- vllm/engine/
|
||||
- vllm/entrypoints/logger.py
|
||||
- vllm/envs.py
|
||||
- vllm/logger.py
|
||||
- vllm/logging_utils/
|
||||
- vllm/platforms/
|
||||
- vllm/sequence.py
|
||||
- vllm/triton_utils/
|
||||
- vllm/utils/
|
||||
- tests/engine
|
||||
- tests/test_sequence
|
||||
- tests/test_config
|
||||
- tests/test_logger
|
||||
- tests/test_vllm_port
|
||||
- tests/test_jit_monitor.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py'
|
||||
|
||||
- label: Engine (1 GPU)
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
@@ -23,3 +61,41 @@ steps:
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
|
||||
|
||||
- label: V1 e2e (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/envs.py
|
||||
- vllm/forward_context.py
|
||||
- vllm/inputs/
|
||||
- vllm/logger.py
|
||||
- vllm/logging_utils/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/triton_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/v1/e2e/spec_decode
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"'
|
||||
|
||||
@@ -125,13 +125,13 @@ steps:
|
||||
pytest -v -s v1/kv_offload &&
|
||||
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
|
||||
|
||||
- label: NixlConnector PD accuracy (2 GPUs)
|
||||
- label: NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 60
|
||||
num_devices: 2
|
||||
num_devices: 4
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
gpu: 4+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
@@ -148,7 +148,10 @@ steps:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
|
||||
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
|
||||
PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
|
||||
PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
|
||||
PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
@@ -259,3 +262,25 @@ steps:
|
||||
pytest -v -s detokenizer &&
|
||||
pytest -v -s -m "not cpu_test" ./multimodal &&
|
||||
pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
|
||||
|
||||
- label: Fusion Unit Tests
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- tests/compile/passes/test_qk_norm_rope_fusion.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
|
||||
@@ -0,0 +1,33 @@
|
||||
group: Model Executor Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: Model Executor (Intel)
|
||||
key: model-executor-intel
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- tests/model_executor
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'apt-get update && apt-get install -y curl libsodium23 &&
|
||||
pip3 install tensorizer==2.10.1 &&
|
||||
pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 &&
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
export PYTHONFAULTHANDLER=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"'
|
||||
@@ -8,7 +8,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -28,7 +28,9 @@ steps:
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
|
||||
pytest -v -s v1/e2e/general/test_context_length.py &&
|
||||
ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
|
||||
pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" &&
|
||||
pytest -v -s v1/e2e/general/test_min_tokens.py'
|
||||
|
||||
- label: Model Runner V2 Examples (Intel)
|
||||
@@ -60,3 +62,55 @@ steps:
|
||||
python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
|
||||
python3 generate/multimodal/vision_language_offline.py --seed 0 &&
|
||||
python3 features/automatic_prefix_caching/prefix_caching_offline.py'
|
||||
|
||||
- label: Model Runner V2 Distributed (2 GPUs)
|
||||
timeout_in_minutes: 50
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/basic_correctness/test_basic_correctness.py
|
||||
- tests/v1/distributed/test_async_llm_dp.py
|
||||
- tests/v1/distributed/test_eagle_dp.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd tests &&
|
||||
TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"'
|
||||
|
||||
- label: Model Runner V2 Spec Decode
|
||||
timeout_in_minutes: 50
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/v1/spec_decode/test_max_len.py
|
||||
- tests/v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- tests/v1/e2e/spec_decode/test_spec_decode.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py'
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
group: Samplers Intel
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: Samplers Test (FlashInfer)
|
||||
key: samplers-test-flashinfer-intel
|
||||
timeout_in_minutes: 40
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers
|
||||
- vllm/sampling_metadata.py
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
- vllm/entrypoints/generate/beam_search
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers'
|
||||
@@ -7,6 +7,9 @@
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
|
||||
git fetch --tags --force origin
|
||||
|
||||
# The Rust frontend build needs protoc.
|
||||
if ! command -v protoc >/dev/null 2>&1; then
|
||||
brew install protobuf
|
||||
|
||||
@@ -35,7 +35,7 @@ set -o pipefail
|
||||
: "${PY_COLORS:=1}"
|
||||
: "${ROCM_DOCKER_TTY:=1}"
|
||||
: "${PYTHONFAULTHANDLER:=1}"
|
||||
: "${PYTEST_TIMEOUT:=2100}"
|
||||
: "${PYTEST_TIMEOUT:=2400}"
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
|
||||
fi
|
||||
@@ -45,9 +45,9 @@ fi
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
|
||||
fi
|
||||
# Dump stacks after 15 minutes, then stop an individual test after 35 minutes.
|
||||
# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900"
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
|
||||
fi
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
|
||||
" ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
|
||||
@@ -387,6 +387,7 @@ initialize_native_environment() {
|
||||
local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
|
||||
local job_id_suffix=""
|
||||
local native_root=""
|
||||
local hf_fstype=""
|
||||
local hf_mount=""
|
||||
|
||||
if [[ "$(id -u)" -ne 0 ]]; then
|
||||
@@ -400,16 +401,19 @@ initialize_native_environment() {
|
||||
native_root="/tmp/vllm-native-${job_id}"
|
||||
TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
|
||||
VLLM_RPC_BASE_PATH="/tmp"
|
||||
: "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}"
|
||||
: "${TRITON_CACHE_DIR:=${native_root}/cache/triton}"
|
||||
: "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}"
|
||||
: "${XDG_CACHE_HOME:=${native_root}/cache/xdg}"
|
||||
TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
|
||||
TRITON_CACHE_DIR="${native_root}/cache/triton"
|
||||
VLLM_CACHE_ROOT="${native_root}/cache/vllm"
|
||||
XDG_CACHE_HOME="${native_root}/cache/xdg"
|
||||
: "${HF_HOME:=/home/buildkite-agent/huggingface}"
|
||||
# datasets uses POSIX locks that are unsupported by the shared HF NFS cache.
|
||||
# Keep processed datasets job-local while retaining the persistent Hub cache.
|
||||
HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets"
|
||||
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
|
||||
: "${HF_HUB_ETAG_TIMEOUT:=60}"
|
||||
export TMPDIR VLLM_RPC_BASE_PATH
|
||||
export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
|
||||
export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
|
||||
export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
|
||||
export PYTORCH_ROCM_ARCH=""
|
||||
|
||||
mkdir -p "${TMPDIR}" \
|
||||
@@ -417,7 +421,10 @@ initialize_native_environment() {
|
||||
"${TRITON_CACHE_DIR}" \
|
||||
"${VLLM_CACHE_ROOT}" \
|
||||
"${XDG_CACHE_HOME}" \
|
||||
"${HF_HOME}" || return 1
|
||||
"${HF_HOME}" \
|
||||
"${HF_DATASETS_CACHE}" || return 1
|
||||
|
||||
echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
|
||||
|
||||
if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
|
||||
if ! command -v findmnt >/dev/null 2>&1; then
|
||||
@@ -430,6 +437,18 @@ initialize_native_environment() {
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if command -v findmnt >/dev/null 2>&1; then
|
||||
hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true)
|
||||
fi
|
||||
if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then
|
||||
# Keep hf-xet state local and avoid vectored writes on shared NFS.
|
||||
export HF_XET_CACHE="${native_root}/cache/hf-xet"
|
||||
export HF_XET_HIGH_PERFORMANCE=0
|
||||
export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0
|
||||
mkdir -p "${HF_XET_CACHE}" || return 1
|
||||
echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}"
|
||||
fi
|
||||
}
|
||||
|
||||
run_native_preflight() {
|
||||
|
||||
+292
-248
@@ -40,7 +40,7 @@
|
||||
#####################################################################################################################################
|
||||
# #
|
||||
# IMPORTANT: #
|
||||
# * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are #
|
||||
# * Currently AMD CI has MI250 agents, MI300 agents, and MI355 agents. All upcoming feature improvements are #
|
||||
# tracked in: https://github.com/vllm-project/vllm/issues/34994 #
|
||||
# #
|
||||
#-----------------------------------------------------------------------------------------------------------------------------------#
|
||||
@@ -81,10 +81,8 @@
|
||||
# the above test.) Also run if model initialization test file is modified. #
|
||||
# * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model #
|
||||
# source is modified, or when specified test files are modified. #
|
||||
# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to #
|
||||
# run plamo2 model in vLLM. #
|
||||
# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) #
|
||||
# and to run plamo2 model in vLLM. #
|
||||
# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). #
|
||||
# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). #
|
||||
# * [Multi-Modal Models (Standard) 1-4]: #
|
||||
# - 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. #
|
||||
@@ -171,20 +169,6 @@ steps:
|
||||
- pip install helion==1.1.0
|
||||
- pytest -v -s kernels/helion/
|
||||
|
||||
- label: Kernels Mamba Test # 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:
|
||||
- csrc/mamba/
|
||||
- tests/kernels/mamba
|
||||
- vllm/model_executor/layers/mamba/ops
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/mamba
|
||||
|
||||
#------------------------------------------------------ mi250 · models / basic -------------------------------------------------------#
|
||||
|
||||
- label: Basic Models Test (Other CPU) # TBD
|
||||
@@ -207,6 +191,7 @@ steps:
|
||||
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/
|
||||
@@ -240,6 +225,20 @@ steps:
|
||||
- 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 Processor (CPU) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
parallelism: 6
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------#
|
||||
|
||||
- label: Batch Invariance (H100-MI250) # TBD
|
||||
@@ -351,21 +350,21 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: V1 attention (H100-MI250) # TBD
|
||||
- label: V1 others (CPU) # 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/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
|
||||
- vllm/
|
||||
- tests/v1
|
||||
commands:
|
||||
- pytest -v -s v1/attention
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
- pytest -v -s v1/structured_output
|
||||
- pytest -v -s v1/test_serial_utils.py
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------#
|
||||
|
||||
@@ -408,6 +407,19 @@ steps:
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
- label: Python-only Installation # 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:
|
||||
- tests/standalone_tests/python_only_compile.sh
|
||||
- setup.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- bash standalone_tests/python_only_compile.sh
|
||||
|
||||
#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------#
|
||||
|
||||
- label: Rust Frontend Cargo Style + Clippy # TBD
|
||||
@@ -445,6 +457,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
no_gpu: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- .buildkite/scripts/docker-build-metadata-args.sh
|
||||
@@ -507,7 +520,7 @@ steps:
|
||||
- tests/models/
|
||||
commands:
|
||||
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
|
||||
- HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
|
||||
- pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
|
||||
@@ -657,6 +670,30 @@ steps:
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
|
||||
|
||||
- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
- vllm/worker/worker_base.py
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/compile/fullgraph/test_basic_correctness.py
|
||||
- tests/compile/test_wrapper.py
|
||||
- tests/entrypoints/llm/test_collective_rpc.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s ./compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s ./compile/test_wrapper.py
|
||||
|
||||
#----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------#
|
||||
|
||||
- label: Platform Tests # TBD
|
||||
@@ -679,6 +716,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -870,6 +908,71 @@ steps:
|
||||
commands:
|
||||
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
|
||||
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
- vllm/worker/worker_base.py
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/distributed/
|
||||
- tests/v1/shutdown
|
||||
- tests/v1/worker/test_worker_memory_snapshot.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
|
||||
- label: Distributed Compile + Comm (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_4
|
||||
num_gpus: 4
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- tests/distributed/test_pynccl
|
||||
- tests/distributed/test_events
|
||||
- tests/compile/fullgraph/test_basic_correctness.py
|
||||
- tests/distributed/test_symm_mem_allreduce.py
|
||||
- tests/distributed/test_multiproc_executor.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s distributed/test_pynccl.py
|
||||
- pytest -v -s distributed/test_events.py
|
||||
- pytest -v -s distributed/test_symm_mem_allreduce.py
|
||||
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
|
||||
#---------------------------------------------------------- mi300 · engine -----------------------------------------------------------#
|
||||
|
||||
- label: Engine # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_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 test_jit_monitor.py
|
||||
|
||||
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
|
||||
|
||||
- label: Entrypoints Unit Tests # TBD
|
||||
@@ -878,6 +981,7 @@ steps:
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
@@ -982,6 +1086,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -1010,6 +1115,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -1024,6 +1130,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -1345,6 +1452,27 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
|
||||
|
||||
- label: LM Eval Large Models (4xH100-4xMI300) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_4
|
||||
num_gpus: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
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
|
||||
commands:
|
||||
- export VLLM_USE_DEEP_GEMM=0
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
|
||||
|
||||
#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------#
|
||||
|
||||
- label: Examples # TBD
|
||||
@@ -1419,11 +1547,12 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: Kernels Core Operation Test # TBD
|
||||
- label: Kernels Core Operation Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
parallelism: 3
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -1434,7 +1563,7 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: Kernels KDA Test # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1452,6 +1581,21 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s kernels/test_kda.py
|
||||
|
||||
- label: Kernels Mamba Test # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/mamba/
|
||||
- tests/kernels/mamba
|
||||
- vllm/model_executor/layers/mamba/ops
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s kernels/mamba
|
||||
|
||||
- label: Kernels MoE Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -1593,7 +1737,7 @@ steps:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
|
||||
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- pytest -v -s v1/e2e/general/test_context_length.py
|
||||
- pytest -v -s v1/e2e/general/test_min_tokens.py
|
||||
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
|
||||
@@ -1797,6 +1941,37 @@ steps:
|
||||
- pip freeze | grep -E 'torch'
|
||||
- pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
- label: Language Models Test (Extended Generation) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
|
||||
- label: Language Models Tests (Hybrid) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
#---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------#
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 1) # TBD
|
||||
@@ -1899,20 +2074,32 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
- label: Multi-Modal Processor (CPU) %N # TBD
|
||||
- label: Multi-Modal Models (Extended Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal/pooling
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/pooling -m 'not core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
parallelism: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
|
||||
#----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------#
|
||||
|
||||
@@ -1933,12 +2120,12 @@ steps:
|
||||
|
||||
#-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------#
|
||||
|
||||
- label: Transformers Nightly Models (Shardable) %N # TBD
|
||||
- label: Transformers Nightly Models (Initialization) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
parallelism: 4
|
||||
parallelism: 6
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
source_file_dependencies:
|
||||
@@ -1954,6 +2141,27 @@ steps:
|
||||
commands:
|
||||
- pip install --upgrade git+https://github.com/huggingface/transformers
|
||||
- pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
- label: Transformers Nightly Models (Processing) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
parallelism: 8
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/multimodal/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
- tests/models/
|
||||
commands:
|
||||
- pip install --upgrade git+https://github.com/huggingface/transformers
|
||||
- pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
- label: Transformers Nightly Models (Single) # TBD
|
||||
@@ -2472,11 +2680,12 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: V1 attention (H100-MI300) # TBD
|
||||
- label: V1 attention (H100-MI300) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -2488,7 +2697,7 @@ steps:
|
||||
- vllm/envs.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s v1/attention
|
||||
- pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: V1 Core + KV + Metrics # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2517,23 +2726,6 @@ steps:
|
||||
# - export HSA_NO_SCRATCH_RECLAIM=1
|
||||
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
|
||||
- label: V1 others (CPU) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1
|
||||
commands:
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
- pytest -v -s v1/structured_output
|
||||
- pytest -v -s v1/test_serial_utils.py
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
- label: V1 Sample + Logits # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -2830,195 +3022,6 @@ steps:
|
||||
commands:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
|
||||
|
||||
#########################################################################################################################################
|
||||
# #
|
||||
# MI325 (gfx942) tests #
|
||||
# #
|
||||
#########################################################################################################################################
|
||||
|
||||
#---------------------------------------------------------- mi325 · compile ----------------------------------------------------------#
|
||||
|
||||
- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
- vllm/worker/worker_base.py
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/compile/fullgraph/test_basic_correctness.py
|
||||
- tests/compile/test_wrapper.py
|
||||
- tests/entrypoints/llm/test_collective_rpc.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s ./compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s ./compile/test_wrapper.py
|
||||
|
||||
#-------------------------------------------------------- mi325 · distributed --------------------------------------------------------#
|
||||
|
||||
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
- vllm/worker/worker_base.py
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/distributed/
|
||||
- tests/v1/shutdown
|
||||
- tests/v1/worker/test_worker_memory_snapshot.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
|
||||
- label: Distributed Compile + Comm (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_4
|
||||
num_gpus: 4
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- tests/distributed/test_pynccl
|
||||
- tests/distributed/test_events
|
||||
- tests/compile/fullgraph/test_basic_correctness.py
|
||||
- tests/distributed/test_symm_mem_allreduce.py
|
||||
- tests/distributed/test_multiproc_executor.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s distributed/test_pynccl.py
|
||||
- pytest -v -s distributed/test_events.py
|
||||
- pytest -v -s distributed/test_symm_mem_allreduce.py
|
||||
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
|
||||
#---------------------------------------------------------- mi325 · engine -----------------------------------------------------------#
|
||||
|
||||
- label: Engine # TBD
|
||||
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/
|
||||
- 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 test_jit_monitor.py
|
||||
|
||||
#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------#
|
||||
|
||||
- label: LM Eval Large Models (4xH100-4xMI325) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_4
|
||||
num_gpus: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
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
|
||||
commands:
|
||||
- export VLLM_USE_DEEP_GEMM=0
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
|
||||
|
||||
#----------------------------------------------------- mi325 · models / language -----------------------------------------------------#
|
||||
|
||||
- label: Language Models Test (Extended Generation) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
|
||||
- label: Language Models Tests (Hybrid) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
parallelism: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
#---------------------------------------------------- mi325 · models / multimodal ----------------------------------------------------#
|
||||
|
||||
- label: Multi-Modal Models (Extended Pooling) # TBD
|
||||
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/
|
||||
- tests/models/multimodal/pooling
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/pooling -m 'not core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
|
||||
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/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
|
||||
#----------------------------------------------------------- mi325 · misc ------------------------------------------------------------#
|
||||
|
||||
- label: Python-only Installation # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- tests/standalone_tests/python_only_compile.sh
|
||||
- setup.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- bash standalone_tests/python_only_compile.sh
|
||||
|
||||
#########################################################################################################################################
|
||||
# #
|
||||
# MI355 (gfx950) tests #
|
||||
@@ -3030,6 +3033,7 @@ steps:
|
||||
- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -3046,6 +3050,7 @@ steps:
|
||||
- label: Distributed Tests (2xH100-2xMI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
@@ -3090,6 +3095,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
@@ -3107,6 +3113,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
optional: true
|
||||
@@ -3122,6 +3129,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
optional: true
|
||||
@@ -3138,6 +3146,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server Generate) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
optional: true
|
||||
@@ -3158,6 +3167,7 @@ steps:
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3171,6 +3181,7 @@ steps:
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3184,6 +3195,7 @@ steps:
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3199,6 +3211,7 @@ steps:
|
||||
- label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
@@ -3221,6 +3234,7 @@ steps:
|
||||
- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
@@ -3243,6 +3257,7 @@ steps:
|
||||
- label: LM Eval Small Models (2xB200-2xMI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
@@ -3262,6 +3277,7 @@ steps:
|
||||
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace"
|
||||
@@ -3282,6 +3298,7 @@ steps:
|
||||
- label: LM Eval Large Models (4xH100-4xMI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_4
|
||||
num_gpus: 4
|
||||
optional: true
|
||||
@@ -3304,6 +3321,7 @@ steps:
|
||||
- label: Examples # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/examples"
|
||||
source_file_dependencies:
|
||||
@@ -3339,6 +3357,7 @@ steps:
|
||||
- label: Kernels (B200-MI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/"
|
||||
source_file_dependencies:
|
||||
@@ -3364,6 +3383,7 @@ steps:
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
parallelism: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3381,6 +3401,7 @@ steps:
|
||||
- label: Kernels MoE Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
parallelism: 5
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3401,6 +3422,7 @@ steps:
|
||||
- label: Kernels Quantization Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
parallelism: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3418,6 +3440,7 @@ steps:
|
||||
- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3437,6 +3460,7 @@ steps:
|
||||
- label: Language Models Test (Extended Generation) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -3450,6 +3474,7 @@ steps:
|
||||
- label: Language Models Test (Extended Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3462,7 +3487,9 @@ steps:
|
||||
- label: Language Models Test (PPL) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/qwen3_5.py
|
||||
@@ -3489,6 +3516,7 @@ steps:
|
||||
- label: Language Models Tests (Standard) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -3503,6 +3531,7 @@ steps:
|
||||
- label: Multi-Modal Models (Extended Generation 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3517,6 +3546,7 @@ steps:
|
||||
- label: Multi-Modal Models (Extended Generation 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3529,6 +3559,7 @@ steps:
|
||||
- label: Multi-Modal Models (Extended Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3541,6 +3572,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3554,6 +3586,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3570,6 +3603,7 @@ steps:
|
||||
- label: Quantized Models Test # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -3586,6 +3620,7 @@ steps:
|
||||
- label: Quantization # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -3602,6 +3637,7 @@ steps:
|
||||
# - label: Quantized MoE Test (B200-MI355) # TBD
|
||||
# timeout_in_minutes: 180
|
||||
# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
# dind: false
|
||||
# agent_pool: mi355_1
|
||||
# working_dir: "/vllm-workspace/"
|
||||
# source_file_dependencies:
|
||||
@@ -3627,10 +3663,12 @@ steps:
|
||||
|
||||
#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------#
|
||||
|
||||
- label: V1 attention (B200-MI355) # TBD
|
||||
- label: V1 attention (B200-MI355) %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
parallelism: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
@@ -3641,11 +3679,12 @@ steps:
|
||||
- vllm/envs.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s v1/attention
|
||||
- pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: V1 Core + KV + Metrics # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3672,6 +3711,7 @@ steps:
|
||||
- label: V1 Sample + Logits # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3692,6 +3732,7 @@ steps:
|
||||
- label: V1 Spec Decode # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -3705,6 +3746,7 @@ steps:
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -3717,6 +3759,7 @@ steps:
|
||||
- label: Weight Loading Multiple GPU - Large Models # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
@@ -3732,6 +3775,7 @@ steps:
|
||||
- label: Regression # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
dind: false
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -16,8 +16,9 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 95
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 125
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -18,7 +18,8 @@ steps:
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 70
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -15,6 +15,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 85
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -68,7 +68,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -94,7 +94,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 85
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -120,7 +120,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 80
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -177,7 +177,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
timeout_in_minutes: 70
|
||||
timeout_in_minutes: 45
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -41,6 +41,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
timeout_in_minutes: 45
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -28,8 +28,9 @@ steps:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 50
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -44,8 +45,8 @@ steps:
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 55
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 45
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -60,8 +61,8 @@ steps:
|
||||
- pytest -v -s v1/e2e/general/test_async_scheduling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 70
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -76,8 +77,8 @@ steps:
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 60
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -116,6 +117,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
timeout_in_minutes: 30
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -30,9 +30,9 @@ steps:
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
# TODO(akaratza): Test after Torch >= 2.12 bump
|
||||
soft_fail: true
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -52,7 +52,9 @@ steps:
|
||||
- pytest -v -s entrypoints/scale_out
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -70,7 +72,8 @@ steps:
|
||||
- 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: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -90,8 +93,9 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 80
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 70
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -113,7 +117,8 @@ steps:
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -176,7 +181,9 @@ steps:
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 30
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -18,6 +18,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 30
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
group: Fault Tolerance
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Fault Tolerance E2E (2xH100)
|
||||
key: fault-tolerance-e2e-2xh100
|
||||
timeout_in_minutes: 35
|
||||
device: h100
|
||||
num_devices: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/fault_tolerance/
|
||||
- vllm/v1/worker/sentinel/
|
||||
- vllm/entrypoints/serve/fault_tolerance/
|
||||
- vllm/distributed/elastic_ep/
|
||||
- vllm/distributed/device_communicators/
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/v1/fault_tolerance/
|
||||
- tests/v1/distributed/test_external_lb_dp.py
|
||||
commands:
|
||||
# Base image has no nixl; install it or has_nixl_ep() skips the tests.
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
|
||||
@@ -80,7 +80,8 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 90
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -118,7 +119,9 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 120
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -148,8 +151,9 @@ steps:
|
||||
parallelism: 5
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 65
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
- csrc/moe/
|
||||
|
||||
@@ -14,8 +14,9 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 55
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 45
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -141,7 +142,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_8
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
@@ -336,7 +337,7 @@ steps:
|
||||
|
||||
- label: LM Eval KV-Offload (2xH100)
|
||||
key: kv-offload-medium
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 45
|
||||
device: h100
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
@@ -346,7 +347,7 @@ steps:
|
||||
- vllm/v1/simple_kv_offload/
|
||||
- tests/evals/gsm8k/test_gsm8k_offloading.py
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
|
||||
|
||||
- label: LM Eval KV-Offload (4xH100)
|
||||
key: kv-offload-large
|
||||
|
||||
@@ -14,9 +14,10 @@ steps:
|
||||
parallelism: 4
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
timeout_in_minutes: 65
|
||||
timeout_in_minutes: 85
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
@@ -46,4 +47,4 @@ steps:
|
||||
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
|
||||
@@ -25,7 +25,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -59,7 +59,9 @@ steps:
|
||||
- pytest -v -s v1/test_outputs.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 70
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -113,8 +115,9 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 75
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -209,7 +212,7 @@ steps:
|
||||
- vllm/multimodal
|
||||
- examples/
|
||||
commands:
|
||||
- pip install tensorizer # for tensorizer test
|
||||
- pip install --no-deps tensorizer # for tensorizer test
|
||||
# for basic
|
||||
- python3 basic/offline_inference/chat.py
|
||||
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
|
||||
@@ -233,7 +236,9 @@ steps:
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- vllm/multimodal
|
||||
@@ -269,10 +274,11 @@ steps:
|
||||
- pytest -v -s v1/tracing
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_2
|
||||
dind: false
|
||||
device: mi300_2
|
||||
timeout_in_minutes: 30
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
optional: true
|
||||
|
||||
- label: Python-only Installation
|
||||
key: python-only-installation
|
||||
@@ -287,8 +293,8 @@ steps:
|
||||
- bash standalone_tests/python_only_compile.sh
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 45
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -30,6 +30,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -41,7 +41,7 @@ steps:
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pip install tensorizer # for tensorizer test
|
||||
- pip install --no-deps tensorizer # for tensorizer test
|
||||
- python3 basic/offline_inference/chat.py # for basic
|
||||
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
|
||||
@@ -42,7 +42,9 @@ steps:
|
||||
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 45
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -39,6 +40,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -61,7 +63,6 @@ steps:
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
# Install fast path packages for testing against transformers
|
||||
# Note: also needed to run plamo2 model in vLLM
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
# Shard the hybrid language model tests that are numerically stable on Hopper.
|
||||
@@ -69,8 +70,9 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 70
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
@@ -102,7 +104,6 @@ steps:
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
# Install fast path packages for testing against transformers
|
||||
# Note: also needed to run plamo2 model in vLLM
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
@@ -130,8 +131,9 @@ steps:
|
||||
- pytest -v -s models/language/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 120
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 95
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -14,7 +14,9 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -31,7 +33,9 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -47,7 +51,8 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -65,7 +70,9 @@ steps:
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -109,6 +116,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 35
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -131,7 +139,9 @@ steps:
|
||||
- pytest -v -s models/multimodal/test_mapping.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 90
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -166,8 +176,9 @@ steps:
|
||||
- pytest -v -s models/multimodal/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 75
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -107,13 +107,6 @@ steps:
|
||||
- tests/compile/passes
|
||||
commands:
|
||||
- pytest -s -v compile/passes --ignore compile/passes/distributed
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: PyTorch Fullgraph Smoke Test
|
||||
device: h200_35gb
|
||||
@@ -236,6 +229,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 30
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -24,8 +24,7 @@ steps:
|
||||
- uv pip install --system conch-triton-kernels
|
||||
# The SM90-only checkpoint currently contains a removed weight_chan_scale
|
||||
# parameter. It was not exercised by the previous L4 job.
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 8
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
|
||||
|
||||
- label: Quantized Fusions
|
||||
device: h200_35gb
|
||||
@@ -68,5 +67,4 @@ steps:
|
||||
- vllm/model_executor/layers/quantization
|
||||
- tests/models/quantization
|
||||
commands:
|
||||
- pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 3
|
||||
- pytest -v -s models/quantization
|
||||
|
||||
@@ -19,8 +19,18 @@ steps:
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers
|
||||
- vllm/sampling_metadata.py
|
||||
- vllm/v1/sample/
|
||||
- vllm/entrypoints/generate/beam_search/
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s samplers
|
||||
|
||||
@@ -14,8 +14,9 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 60
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -53,8 +54,9 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 65
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 75
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -92,10 +94,9 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 55
|
||||
# TODO(akaratza): Test after Torch >= 2.12 bump
|
||||
soft_fail: true
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 35
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
@@ -119,7 +120,8 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -17,6 +17,7 @@ steps:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
timeout_in_minutes: 35
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
|
||||
@@ -19,6 +19,7 @@ pull_request_rules:
|
||||
description: Comment on PR when pre-commit check fails
|
||||
conditions:
|
||||
- check-failure=pre-commit
|
||||
- -check-cancelled=pre-commit
|
||||
- -closed
|
||||
- -draft
|
||||
- or:
|
||||
@@ -232,6 +233,31 @@ pull_request_rules:
|
||||
add:
|
||||
- gpt-oss
|
||||
|
||||
- name: label-kimi
|
||||
description: Automatically apply kimi label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=(?i)kimi
|
||||
- files~=(?i)moonshot
|
||||
- title~=(?i)(?:kimi|moonshot)
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- kimi
|
||||
|
||||
- name: label-k3
|
||||
description: Automatically apply k3 label (launch triage; retire after ramp-down)
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=(?i)kimi[-_]?k3
|
||||
- title~=(?i)(?:kimi[-\s]?k3|\bk3\b)
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- k3
|
||||
|
||||
- name: label-nvidia
|
||||
description: Automatically apply nvidia label
|
||||
conditions:
|
||||
|
||||
@@ -130,6 +130,25 @@ jobs:
|
||||
},
|
||||
],
|
||||
},
|
||||
kimi: {
|
||||
keywords: [
|
||||
{ term: "Kimi", searchIn: "both" },
|
||||
{ term: "Moonshot", searchIn: "both" },
|
||||
],
|
||||
substrings: [
|
||||
{ term: "moonshotai/", searchIn: "both" },
|
||||
{ term: "kimi", searchIn: "title" },
|
||||
],
|
||||
},
|
||||
k3: {
|
||||
keywords: [
|
||||
{ term: "Kimi K3", searchIn: "both" },
|
||||
{ term: "K3", searchIn: "title" },
|
||||
],
|
||||
substrings: [
|
||||
{ term: "moonshotai/kimi-k3", searchIn: "both" },
|
||||
],
|
||||
},
|
||||
quantization: {
|
||||
keywords: [
|
||||
{
|
||||
|
||||
@@ -173,9 +173,6 @@ venv.bak/
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
docs/argparse
|
||||
docs/examples/*
|
||||
!docs/examples/README.md
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
|
||||
@@ -3,6 +3,9 @@ MD007:
|
||||
MD013: false
|
||||
MD024:
|
||||
siblings_only: true
|
||||
MD025:
|
||||
# Allow front matter title to be different from the first heading in the document.
|
||||
front_matter_title: ""
|
||||
MD031:
|
||||
list_items: false
|
||||
MD033: false
|
||||
|
||||
@@ -260,10 +260,6 @@ repos:
|
||||
files: ^docker/(Dockerfile|versions\.json)$
|
||||
pass_filenames: false
|
||||
additional_dependencies: [dockerfile-parse]
|
||||
- id: attention-backend-docs
|
||||
name: Check attention backend documentation is up to date
|
||||
entry: python tools/pre_commit/generate_attention_backend_docs.py --check
|
||||
language: python
|
||||
- id: check-boolean-context-manager
|
||||
name: Check for boolean ops in with-statements
|
||||
entry: python tools/pre_commit/check_boolean_context_manager.py
|
||||
|
||||
+12
-7
@@ -114,6 +114,11 @@ find_package(Torch REQUIRED)
|
||||
# Supported NVIDIA architectures.
|
||||
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
|
||||
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
|
||||
# Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
|
||||
# targeting it directly.
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
|
||||
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
|
||||
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
|
||||
# to support the whole generation without specifying all sub-architectures
|
||||
@@ -420,7 +425,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
|
||||
"9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
|
||||
"9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
|
||||
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
@@ -695,7 +700,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
@@ -815,7 +820,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
|
||||
# require CUDA 12.8 or later
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
@@ -899,7 +904,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
@@ -924,7 +929,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
# moe_data.cu is used by all CUTLASS MoE kernels.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
@@ -981,7 +986,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
|
||||
# only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
@@ -1047,7 +1052,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Runtime dispatch is gated in
|
||||
# vllm/v1/attention/backends/mla/cutlass_mla.py.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
@@ -1358,6 +1358,10 @@ def main():
|
||||
profile_memory=args.profile_memory,
|
||||
warmup_ms=args.warmup_ms,
|
||||
prefill_backend=pb,
|
||||
kv_lora_rank=args.kv_lora_rank,
|
||||
qk_nope_head_dim=args.qk_nope_head_dim,
|
||||
qk_rope_head_dim=args.qk_rope_head_dim,
|
||||
v_head_dim=args.v_head_dim,
|
||||
)
|
||||
|
||||
result = run_benchmark(config)
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import statistics
|
||||
|
||||
import torch
|
||||
from tabulate import tabulate
|
||||
|
||||
from vllm.models.inkling.nvidia.ops import qkvr_prep
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def make_inputs(tokens: int, tp_size: int, is_local: bool):
|
||||
torch.manual_seed(0)
|
||||
num_q_heads = 64 // tp_size
|
||||
num_kv_heads = (16 if is_local else 8) // tp_size
|
||||
head_dim = 128
|
||||
d_rel = 16
|
||||
rel_extent = 512 if is_local else 1024
|
||||
page_size = 16
|
||||
num_blocks = (tokens + page_size - 1) // page_size
|
||||
q_width = num_q_heads * head_dim
|
||||
kv_width = num_kv_heads * head_dim
|
||||
r_width = num_q_heads * d_rel
|
||||
device = "cuda"
|
||||
|
||||
qkvr = torch.randn(
|
||||
tokens,
|
||||
q_width + 2 * kv_width + r_width,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
|
||||
v_weight = torch.randn_like(k_weight)
|
||||
q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
|
||||
k_norm_weight = torch.randn_like(q_norm_weight)
|
||||
rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
|
||||
conv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
num_kv_heads,
|
||||
page_size,
|
||||
2 * head_dim,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
key_cache = torch.empty(
|
||||
num_blocks,
|
||||
page_size,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
value_cache = torch.empty_like(key_cache)
|
||||
positions = torch.arange(tokens, device=device, dtype=torch.int64)
|
||||
block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
|
||||
seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
|
||||
slots = torch.arange(tokens, device=device, dtype=torch.int64)
|
||||
query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
|
||||
log_scaling = None
|
||||
if not is_local:
|
||||
effective_n = (positions + 1).to(torch.float32)
|
||||
log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0))
|
||||
return (
|
||||
qkvr,
|
||||
k_weight,
|
||||
v_weight,
|
||||
q_norm_weight,
|
||||
k_norm_weight,
|
||||
rel_proj,
|
||||
1e-6,
|
||||
num_q_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
d_rel,
|
||||
conv_cache,
|
||||
key_cache,
|
||||
value_cache,
|
||||
positions,
|
||||
block_table,
|
||||
seq_idx,
|
||||
slots,
|
||||
query_start,
|
||||
slots,
|
||||
0,
|
||||
head_dim,
|
||||
page_size,
|
||||
log_scaling,
|
||||
)
|
||||
|
||||
|
||||
def capture(implementation, inputs):
|
||||
outputs = []
|
||||
|
||||
def run():
|
||||
outputs[:] = implementation.fused_qkvr_prep(*inputs)
|
||||
|
||||
stream = torch.cuda.Stream()
|
||||
stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(stream):
|
||||
for _ in range(3):
|
||||
run()
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
torch.accelerator.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
run()
|
||||
torch.accelerator.synchronize()
|
||||
return graph, outputs
|
||||
|
||||
|
||||
def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float:
|
||||
for _ in range(warmup):
|
||||
graph.replay()
|
||||
torch.accelerator.synchronize()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(repeats):
|
||||
graph.replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) * 1000 / repeats
|
||||
|
||||
|
||||
def benchmark(inputs, args) -> float:
|
||||
graph, _ = capture(qkvr_prep, inputs)
|
||||
return statistics.median(
|
||||
time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials)
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main(args):
|
||||
rows = []
|
||||
for tp_size in args.tp_sizes:
|
||||
for tokens in args.tokens:
|
||||
for is_local in (True, False):
|
||||
triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args)
|
||||
rows.append(
|
||||
[
|
||||
tp_size,
|
||||
tokens,
|
||||
"local" if is_local else "global",
|
||||
triton_us,
|
||||
]
|
||||
)
|
||||
|
||||
print("Inkling QKVR prep (CUDA graph, median latency)")
|
||||
print(
|
||||
tabulate(
|
||||
rows,
|
||||
headers=[
|
||||
"TP",
|
||||
"tokens",
|
||||
"scope",
|
||||
"Triton (us)",
|
||||
],
|
||||
floatfmt=("d", "d", "", ".2f"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser()
|
||||
parser.add_argument(
|
||||
"--tokens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1 << power for power in range(15)],
|
||||
)
|
||||
parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8])
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--repeats", type=int, default=200)
|
||||
parser.add_argument("--trials", type=int, default=5)
|
||||
main(parser.parse_args())
|
||||
@@ -154,7 +154,7 @@ def main(
|
||||
scale=scale,
|
||||
causal=True,
|
||||
alibi_slopes=None,
|
||||
sliding_window=window_size,
|
||||
sliding_window=window_size if sliding_window is not None else -1,
|
||||
block_table=block_tables,
|
||||
softcap=0,
|
||||
scheduler_metadata=metadata,
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
|
||||
|
||||
Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
|
||||
long greedy decode (CUDA graphs on) once with the standard kernel and once with
|
||||
ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
|
||||
separate subprocesses so each gets a clean CUDA context.
|
||||
|
||||
The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
|
||||
CUDA-graph capture on the pre-release Blackwell FP4 path); pass
|
||||
--no-disable-flashinfer-autotune for non-FP4 models.
|
||||
|
||||
Examples:
|
||||
python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
|
||||
python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
|
||||
--model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
|
||||
|
||||
MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(
|
||||
description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
|
||||
)
|
||||
p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
|
||||
p.add_argument("--prompt", default=DEFAULT_PROMPT)
|
||||
p.add_argument("--batch-size", type=int, default=256)
|
||||
p.add_argument("--num-steps", type=int, default=1000)
|
||||
p.add_argument("--warmup-steps", type=int, default=128)
|
||||
p.add_argument("--repeats", type=int, default=1)
|
||||
p.add_argument(
|
||||
"--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
|
||||
)
|
||||
p.add_argument(
|
||||
"--dtype",
|
||||
default="bfloat16",
|
||||
choices=["bfloat16", "float16", "float32", "auto"],
|
||||
)
|
||||
p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
|
||||
p.add_argument("--max-model-len", type=int, default=None)
|
||||
p.add_argument(
|
||||
"--disable-flashinfer-autotune",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
|
||||
"It is unstable under CUDA-graph capture on the "
|
||||
"pre-release Blackwell FP4 path; pass "
|
||||
"--no-disable-flashinfer-autotune for non-FP4 models.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--mamba-ssm-cache-dtype",
|
||||
default="auto",
|
||||
choices=["auto", "float32", "float16", "bfloat16"],
|
||||
help="SSM state dtype (both modes). 'auto' = config-driven; "
|
||||
"'float32' = fp32 state, 'bfloat16' = s16 state.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--baseline-ssm-config",
|
||||
default="",
|
||||
help="Pin the STANDARD baseline's SSM launch config as "
|
||||
"'bsm,nw' via override_ssm_config (forces the in-process "
|
||||
"engine so the override reaches the kernel). Empty = off.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--worker",
|
||||
choices=["standard", "replayssm"],
|
||||
default=None,
|
||||
help=argparse.SUPPRESS,
|
||||
)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def resolve_max_model_len(args) -> int:
|
||||
if args.max_model_len is not None:
|
||||
return args.max_model_len
|
||||
return args.num_steps + 256
|
||||
|
||||
|
||||
def run_worker(args):
|
||||
# override_ssm_config is a module global; it only reaches the model if the
|
||||
# engine runs in-process (default V1 spawns a separate EngineCore). Force it.
|
||||
if args.worker == "standard" and args.baseline_ssm_config:
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
mode = args.worker
|
||||
max_model_len = resolve_max_model_len(args)
|
||||
|
||||
llm_kwargs = dict(
|
||||
model=args.model_id,
|
||||
tensor_parallel_size=1,
|
||||
dtype=args.dtype,
|
||||
max_model_len=max_model_len,
|
||||
trust_remote_code=True,
|
||||
enable_prefix_caching=False,
|
||||
enable_chunked_prefill=False,
|
||||
max_num_seqs=args.batch_size,
|
||||
max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
|
||||
enforce_eager=False,
|
||||
disable_log_stats=True,
|
||||
gpu_memory_utilization=args.gpu_memory_utilization,
|
||||
# SSM state dtype (applies to both standard and ReplaySSM).
|
||||
mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
|
||||
)
|
||||
if args.disable_flashinfer_autotune:
|
||||
# FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
|
||||
# re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
|
||||
llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
|
||||
if mode == "replayssm":
|
||||
llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
|
||||
|
||||
_ssm_cm = None
|
||||
if mode == "standard" and args.baseline_ssm_config:
|
||||
from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
|
||||
|
||||
_bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
|
||||
_ssm_cm = override_ssm_config((_bsm, _nw))
|
||||
_ssm_cm.__enter__() # active through LLM() graph capture + decode
|
||||
print(
|
||||
f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
llm = LLM(**llm_kwargs)
|
||||
prompts = [args.prompt] * args.batch_size
|
||||
|
||||
def timed_generate(n_tokens):
|
||||
sp = SamplingParams(
|
||||
n=1,
|
||||
temperature=0.0,
|
||||
ignore_eos=True,
|
||||
min_tokens=n_tokens,
|
||||
max_tokens=n_tokens,
|
||||
)
|
||||
if torch.accelerator.is_available():
|
||||
torch.accelerator.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
outs = llm.generate(prompts, sp, use_tqdm=False)
|
||||
if torch.accelerator.is_available():
|
||||
torch.accelerator.synchronize()
|
||||
elapsed = time.perf_counter() - t0
|
||||
produced = min(len(o.outputs[0].token_ids) for o in outs)
|
||||
assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
|
||||
return elapsed
|
||||
|
||||
timed_generate(args.warmup_steps)
|
||||
|
||||
best = None
|
||||
for _ in range(args.repeats):
|
||||
elapsed = timed_generate(args.num_steps)
|
||||
tok_s = args.batch_size * args.num_steps / elapsed
|
||||
per_step_ms = elapsed / args.num_steps * 1e3
|
||||
print(
|
||||
f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step",
|
||||
flush=True,
|
||||
)
|
||||
if best is None or elapsed < best["elapsed_s"]:
|
||||
best = {
|
||||
"mode": mode,
|
||||
"elapsed_s": elapsed,
|
||||
"tok_s": tok_s,
|
||||
"per_step_ms": per_step_ms,
|
||||
}
|
||||
|
||||
print("RESULT_JSON " + json.dumps(best), flush=True)
|
||||
if _ssm_cm is not None:
|
||||
_ssm_cm.__exit__(None, None, None)
|
||||
|
||||
|
||||
def run_one_mode(args, mode) -> dict:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
__file__,
|
||||
"--worker",
|
||||
mode,
|
||||
"--model-id",
|
||||
args.model_id,
|
||||
"--prompt",
|
||||
args.prompt,
|
||||
"--batch-size",
|
||||
str(args.batch_size),
|
||||
"--num-steps",
|
||||
str(args.num_steps),
|
||||
"--warmup-steps",
|
||||
str(args.warmup_steps),
|
||||
"--repeats",
|
||||
str(args.repeats),
|
||||
"--buffer-len",
|
||||
str(args.buffer_len),
|
||||
"--dtype",
|
||||
args.dtype,
|
||||
"--gpu-memory-utilization",
|
||||
str(args.gpu_memory_utilization),
|
||||
"--mamba-ssm-cache-dtype",
|
||||
args.mamba_ssm_cache_dtype,
|
||||
"--baseline-ssm-config",
|
||||
args.baseline_ssm_config,
|
||||
]
|
||||
cmd.append(
|
||||
"--disable-flashinfer-autotune"
|
||||
if args.disable_flashinfer_autotune
|
||||
else "--no-disable-flashinfer-autotune"
|
||||
)
|
||||
if args.max_model_len is not None:
|
||||
cmd += ["--max-model-len", str(args.max_model_len)]
|
||||
|
||||
result = None
|
||||
proc = subprocess.Popen(
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
|
||||
)
|
||||
for line in proc.stdout:
|
||||
sys.stdout.write(line)
|
||||
sys.stdout.flush()
|
||||
if line.startswith("RESULT_JSON "):
|
||||
result = json.loads(line[len("RESULT_JSON ") :])
|
||||
proc.wait()
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
|
||||
if result is None:
|
||||
raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
|
||||
return result
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
if args.worker is not None:
|
||||
run_worker(args)
|
||||
return
|
||||
|
||||
print(
|
||||
f"model={args.model_id} batch_size={args.batch_size} "
|
||||
f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}"
|
||||
)
|
||||
|
||||
std = run_one_mode(args, "standard")
|
||||
fla = run_one_mode(args, "replayssm")
|
||||
speedup = std["per_step_ms"] / fla["per_step_ms"]
|
||||
|
||||
print()
|
||||
header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for r in (std, fla):
|
||||
print(
|
||||
f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
|
||||
f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
|
||||
)
|
||||
print("-" * len(header))
|
||||
print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -15,6 +15,7 @@ endif()
|
||||
#
|
||||
set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
|
||||
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
|
||||
set(ENABLE_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM})
|
||||
set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
|
||||
|
||||
include_directories("${CMAKE_SOURCE_DIR}/csrc")
|
||||
@@ -96,12 +97,14 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
|
||||
set(ENABLE_NUMA OFF)
|
||||
check_sysctl(hw.optional.neon ASIMD_FOUND)
|
||||
check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
|
||||
check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND)
|
||||
else()
|
||||
find_isa(${CPUINFO} "Power11" POWER11_FOUND)
|
||||
find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
|
||||
find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
|
||||
find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
|
||||
find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
|
||||
find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support
|
||||
find_isa(${CPUINFO} "S390" S390_FOUND)
|
||||
find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
|
||||
find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
|
||||
@@ -111,6 +114,11 @@ else()
|
||||
set(ARM_BF16_FOUND ON)
|
||||
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
|
||||
endif()
|
||||
if (ENABLE_ARM_I8MM)
|
||||
set(ARM_I8MM_FOUND ON)
|
||||
message(STATUS
|
||||
"ARM I8MM support enabled via VLLM_CPU_ARM_I8MM environment variable")
|
||||
endif()
|
||||
# Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
|
||||
# in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
|
||||
# overrides the detection result.
|
||||
@@ -166,6 +174,11 @@ elseif (ASIMD_FOUND)
|
||||
message(WARNING "BF16 functionality is not available")
|
||||
set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")
|
||||
endif()
|
||||
if(ARM_I8MM_FOUND)
|
||||
message(STATUS "I8MM extension detected")
|
||||
string(APPEND MARCH_FLAGS "+i8mm")
|
||||
add_compile_definitions(ARM_I8MM_SUPPORT)
|
||||
endif()
|
||||
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
|
||||
elseif (S390_FOUND)
|
||||
message(STATUS "S390 detected")
|
||||
@@ -447,8 +460,13 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
"csrc/cpu/shm.cpp"
|
||||
"csrc/cpu/activation_lut_bf16.cpp"
|
||||
"csrc/cpu/cpu_tanhf_neon.hpp"
|
||||
"csrc/cpu/cpu_fused_moe.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
if (ARM_BF16_FOUND)
|
||||
set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC})
|
||||
if (ARM_I8MM_FOUND)
|
||||
set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
|
||||
endif()
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
|
||||
|
||||
@@ -68,6 +68,9 @@ endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
|
||||
endif()
|
||||
else()
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
|
||||
endif()
|
||||
|
||||
@@ -60,6 +60,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
|
||||
# CUDA 12.9 has introduced "Family-Specific Architecture Features"
|
||||
# this supports all compute_10x family
|
||||
list(APPEND SUPPORT_ARCHS "10.0f")
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
|
||||
list(APPEND SUPPORT_ARCHS "10.7f")
|
||||
endif()
|
||||
elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
|
||||
list(APPEND SUPPORT_ARCHS "10.0a")
|
||||
endif()
|
||||
@@ -188,4 +191,3 @@ else()
|
||||
add_custom_target(_flashmla_C)
|
||||
add_custom_target(_flashmla_extension_C)
|
||||
endif()
|
||||
|
||||
|
||||
@@ -55,7 +55,11 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
|
||||
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
else()
|
||||
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
|
||||
|
||||
+15
-5
@@ -396,14 +396,24 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
|
||||
# match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
|
||||
# family. The output uses TGT's value to preserve the user's compilation flags.
|
||||
set(_CUDA_ARCHS)
|
||||
# Resolve exact base matches before family fallbacks so a generic entry such
|
||||
# as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
|
||||
foreach(_arch ${_SRC_CUDA_ARCHS})
|
||||
if(_arch MATCHES "[af]$")
|
||||
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
|
||||
if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
|
||||
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
|
||||
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
|
||||
list(APPEND _CUDA_ARCHS "${_arch}")
|
||||
endif()
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
foreach(_arch ${_SRC_CUDA_ARCHS})
|
||||
if(_arch MATCHES "[af]$")
|
||||
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
|
||||
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
|
||||
if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
|
||||
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
|
||||
list(APPEND _CUDA_ARCHS "${_arch}")
|
||||
elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
|
||||
if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
|
||||
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
|
||||
list(APPEND _CUDA_ARCHS "${_base}a")
|
||||
elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
|
||||
@@ -487,7 +497,7 @@ endfunction()
|
||||
|
||||
function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
@@ -172,4 +172,15 @@
|
||||
|
||||
#endif // __riscv_v
|
||||
|
||||
// Power VSX
|
||||
#ifdef __powerpc__
|
||||
// FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which
|
||||
// implements a vectorised 5-term minimax polynomial using VSX intrinsics.
|
||||
#define DEFINE_FAST_EXP \
|
||||
auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
|
||||
__attribute__((always_inline)) { return vec.exp(); }; \
|
||||
auto fast_exp_f16 = fast_exp;
|
||||
|
||||
#endif // __powerpc__
|
||||
|
||||
#endif
|
||||
|
||||
+5
-187
@@ -1,5 +1,6 @@
|
||||
#include "cpu/cpu_types.hpp"
|
||||
#include "cpu/utils.hpp"
|
||||
#include "cpu/cpu_fused_moe_activations.hpp"
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
|
||||
#include "cpu/cpu_arch_macros.h"
|
||||
|
||||
@@ -43,193 +44,9 @@
|
||||
}()
|
||||
|
||||
namespace {
|
||||
enum class FusedMOEAct {
|
||||
SiluAndMul,
|
||||
SwigluOAIAndMul,
|
||||
GeluAndMul,
|
||||
GeluTanhAndMul,
|
||||
};
|
||||
|
||||
FusedMOEAct get_act_type(const std::string& act) {
|
||||
if (act == "silu") {
|
||||
return FusedMOEAct::SiluAndMul;
|
||||
} else if (act == "swigluoai") {
|
||||
return FusedMOEAct::SwigluOAIAndMul;
|
||||
} else if (act == "gelu") {
|
||||
return FusedMOEAct::GeluAndMul;
|
||||
} else if (act == "gelu_tanh") {
|
||||
return FusedMOEAct::GeluTanhAndMul;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid act type: " + act);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
#if !defined(__aarch64__)
|
||||
// For GPT-OSS interleaved gate-up weights
|
||||
alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
|
||||
16, 18, 20, 22, 24, 26, 28, 30};
|
||||
vec_op::INT32Vec16 index_vec(index);
|
||||
#endif
|
||||
vec_op::FP32Vec16 gate_up_max_vec(7.0);
|
||||
vec_op::FP32Vec16 up_min_vec(-7.0);
|
||||
vec_op::FP32Vec16 alpha_vec(1.702);
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < n_size; n += 32) {
|
||||
// Note: AdvSIMD does not support gather loads
|
||||
#if defined(__aarch64__)
|
||||
vec_op::FP32Vec16 gate_vec(vec_op::uninit);
|
||||
vec_op::FP32Vec16 up_vec(vec_op::uninit);
|
||||
vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
|
||||
#else
|
||||
vec_op::FP32Vec16 gate_vec(input + n, index_vec);
|
||||
vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
|
||||
#endif
|
||||
gate_vec = gate_vec.min(gate_up_max_vec);
|
||||
up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
|
||||
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
|
||||
auto glu = gate_vec * sigmoid_vec;
|
||||
auto gated_output_fp32 = (one_vec + up_vec) * glu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n / 2);
|
||||
}
|
||||
input += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride, const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
|
||||
auto silu = gate_vec * sigmoid_vec;
|
||||
auto gated_output_fp32 = up_vec * silu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride, const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
|
||||
vec_op::FP32Vec16 w2_vec(0.5);
|
||||
alignas(64) float temp[16];
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto er_input_vec = gate_vec * w1_vec;
|
||||
|
||||
er_input_vec.save(temp);
|
||||
for (int32_t i = 0; i < 16; ++i) {
|
||||
temp[i] = std::erf(temp[i]);
|
||||
}
|
||||
vec_op::FP32Vec16 er_vec(temp);
|
||||
auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
|
||||
auto gated_output_fp32 = up_vec * gelu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
|
||||
vec_op::FP32Vec16 w2_vec(0.5);
|
||||
vec_op::FP32Vec16 w3_vec(0.044715);
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
|
||||
auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
|
||||
// Note: can't use fast_exp form because diffusiongemma will generate
|
||||
// wrong results
|
||||
auto tanh_vec = inner_vec.tanh();
|
||||
auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
|
||||
auto gated_output_fp32 = up_vec * gelu_tanh;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
|
||||
float* __restrict__ input,
|
||||
scalar_t* __restrict__ output,
|
||||
const int32_t m, const int32_t n,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
switch (act) {
|
||||
case FusedMOEAct::SwigluOAIAndMul:
|
||||
swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::SiluAndMul:
|
||||
silu_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::GeluAndMul:
|
||||
gelu_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::GeluTanhAndMul:
|
||||
gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unsupported act type.");
|
||||
}
|
||||
}
|
||||
using cpu_fused_moe_utils::apply_gated_act;
|
||||
using cpu_fused_moe_utils::FusedMOEAct;
|
||||
|
||||
template <typename scalar_t, typename gemm_t>
|
||||
void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
|
||||
@@ -817,6 +634,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
void prepack_moe_weight(
|
||||
@@ -864,7 +682,7 @@ void cpu_fused_moe(
|
||||
const int32_t input_size_2 = w2.size(2);
|
||||
const int32_t output_size_2 = w2.size(1);
|
||||
const int32_t topk_num = topk_id.size(1);
|
||||
const FusedMOEAct act_type = get_act_type(act);
|
||||
const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
|
||||
cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
|
||||
TORCH_CHECK(!skip_weighted || topk_num == 1,
|
||||
"skip_weighted is only supported for topk=1 on CPU");
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP
|
||||
#define CPU_FUSED_MOE_ACTIVATIONS_HPP
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
|
||||
#include "cpu/cpu_arch_macros.h"
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
namespace cpu_fused_moe_utils {
|
||||
enum class FusedMOEAct {
|
||||
SiluAndMul,
|
||||
SwigluOAIAndMul,
|
||||
GeluAndMul,
|
||||
GeluTanhAndMul,
|
||||
};
|
||||
|
||||
inline FusedMOEAct get_act_type(const std::string& act) {
|
||||
if (act == "silu") {
|
||||
return FusedMOEAct::SiluAndMul;
|
||||
} else if (act == "swigluoai") {
|
||||
return FusedMOEAct::SwigluOAIAndMul;
|
||||
} else if (act == "gelu") {
|
||||
return FusedMOEAct::GeluAndMul;
|
||||
} else if (act == "gelu_tanh") {
|
||||
return FusedMOEAct::GeluTanhAndMul;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid act type: " + act);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
#if !defined(__aarch64__)
|
||||
// For GPT-OSS interleaved gate-up weights
|
||||
alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
|
||||
16, 18, 20, 22, 24, 26, 28, 30};
|
||||
vec_op::INT32Vec16 index_vec(index);
|
||||
#endif
|
||||
vec_op::FP32Vec16 gate_up_max_vec(7.0);
|
||||
vec_op::FP32Vec16 up_min_vec(-7.0);
|
||||
vec_op::FP32Vec16 alpha_vec(1.702);
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < n_size; n += 32) {
|
||||
// Note: AdvSIMD does not support gather loads
|
||||
#if defined(__aarch64__)
|
||||
vec_op::FP32Vec16 gate_vec(vec_op::uninit);
|
||||
vec_op::FP32Vec16 up_vec(vec_op::uninit);
|
||||
vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
|
||||
#else
|
||||
vec_op::FP32Vec16 gate_vec(input + n, index_vec);
|
||||
vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
|
||||
#endif
|
||||
gate_vec = gate_vec.min(gate_up_max_vec);
|
||||
up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
|
||||
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
|
||||
auto glu = gate_vec * sigmoid_vec;
|
||||
auto gated_output_fp32 = (one_vec + up_vec) * glu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n / 2);
|
||||
}
|
||||
input += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride, const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
|
||||
auto silu = gate_vec * sigmoid_vec;
|
||||
auto gated_output_fp32 = up_vec * silu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride, const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
|
||||
vec_op::FP32Vec16 w2_vec(0.5);
|
||||
alignas(64) float temp[16];
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto er_input_vec = gate_vec * w1_vec;
|
||||
|
||||
er_input_vec.save(temp);
|
||||
for (int32_t i = 0; i < 16; ++i) {
|
||||
temp[i] = std::erf(temp[i]);
|
||||
}
|
||||
vec_op::FP32Vec16 er_vec(temp);
|
||||
auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
|
||||
auto gated_output_fp32 = up_vec * gelu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
|
||||
vec_op::FP32Vec16 w2_vec(0.5);
|
||||
vec_op::FP32Vec16 w3_vec(0.044715);
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
|
||||
auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
|
||||
// Note: can't use fast_exp form because diffusiongemma will generate
|
||||
// wrong results
|
||||
auto tanh_vec = inner_vec.tanh();
|
||||
auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
|
||||
auto gated_output_fp32 = up_vec * gelu_tanh;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
|
||||
float* __restrict__ input,
|
||||
scalar_t* __restrict__ output,
|
||||
const int32_t m, const int32_t n,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
switch (act) {
|
||||
case FusedMOEAct::SwigluOAIAndMul:
|
||||
swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::SiluAndMul:
|
||||
silu_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::GeluAndMul:
|
||||
gelu_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::GeluTanhAndMul:
|
||||
gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unsupported act type.");
|
||||
}
|
||||
}
|
||||
} // namespace cpu_fused_moe_utils
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,647 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include "cpu/cpu_arch_macros.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
|
||||
#include "cpu/cpu_fused_moe_activations.hpp"
|
||||
#include "cpu/cpu_types.hpp"
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT)
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp"
|
||||
#define NEON_DISPATCH(SCALAR_TYPE, ...) \
|
||||
case cpu_utils::ISA::NEON: { \
|
||||
using gemm_t = \
|
||||
cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::NEON, SCALAR_TYPE>; \
|
||||
return __VA_ARGS__(); \
|
||||
}
|
||||
#else
|
||||
#define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON:
|
||||
#endif
|
||||
|
||||
#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \
|
||||
[&] { \
|
||||
switch (ISA_TYPE) { \
|
||||
NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \
|
||||
default: { \
|
||||
TORCH_CHECK(false, "Invalid CPU ISA type."); \
|
||||
} \
|
||||
} \
|
||||
}()
|
||||
|
||||
namespace {
|
||||
using cpu_fused_moe_utils::apply_gated_act;
|
||||
using cpu_fused_moe_utils::FusedMOEAct;
|
||||
|
||||
template <typename gemm_t>
|
||||
void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr,
|
||||
int8_t* __restrict__ packed_weight_ptr,
|
||||
const int32_t expert_num,
|
||||
const int32_t output_size,
|
||||
const int32_t input_size,
|
||||
const int64_t expert_stride) {
|
||||
#pragma omp parallel for
|
||||
for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) {
|
||||
gemm_t::pack_weight(weight_ptr + expert_stride * e_idx,
|
||||
packed_weight_ptr + expert_stride * e_idx, output_size,
|
||||
input_size);
|
||||
}
|
||||
}
|
||||
|
||||
// INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp
|
||||
template <typename scalar_t, typename gemm_t>
|
||||
void fused_moe_int8_impl(
|
||||
scalar_t* __restrict__ output, const scalar_t* __restrict__ input,
|
||||
const int8_t* __restrict__ w13, const int8_t* __restrict__ w2,
|
||||
const float* __restrict__ w13_scales, const float* __restrict__ w2_scales,
|
||||
scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias,
|
||||
const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id,
|
||||
const FusedMOEAct act_type, const int32_t token_num,
|
||||
const int32_t expert_num, const int32_t topk_num,
|
||||
const int32_t input_size_13, const int32_t output_size_13,
|
||||
const int32_t input_size_2, const int32_t output_size_2,
|
||||
const bool skip_weighted) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
|
||||
constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
|
||||
constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
|
||||
|
||||
TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0);
|
||||
TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 0);
|
||||
TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
|
||||
TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0);
|
||||
TORCH_CHECK_EQ(output_size_13 / 2, input_size_2);
|
||||
|
||||
const int32_t thread_num = cpu_utils::get_max_threads();
|
||||
const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
|
||||
gemm_m_tile_size * input_size_13 * sizeof(int8_t));
|
||||
const int32_t w2_input_buffer_size =
|
||||
cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t));
|
||||
|
||||
const int32_t w13_n_tile_size = [&]() {
|
||||
const int64_t cache_size = cpu_utils::get_available_l2_size();
|
||||
const int32_t n_size_cache_limit =
|
||||
(cache_size - w13_input_buffer_size) /
|
||||
(gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t));
|
||||
const int32_t n_size_thread_limit =
|
||||
output_size_13 / std::max(1, thread_num / topk_num);
|
||||
const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>(
|
||||
std::min(n_size_cache_limit, n_size_thread_limit));
|
||||
return std::max(n_size, min_w13_n_tile_size);
|
||||
}();
|
||||
|
||||
const int32_t w2_n_tile_size = [&]() {
|
||||
const int64_t cache_size = cpu_utils::get_available_l2_size();
|
||||
const int32_t n_size_cache_limit =
|
||||
(cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_t));
|
||||
const int32_t n_size_thread_limit =
|
||||
output_size_2 / std::max(1, thread_num / topk_num);
|
||||
const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
|
||||
std::min(n_size_cache_limit, n_size_thread_limit));
|
||||
return std::max(n_size, gemm_n_tile_size);
|
||||
}();
|
||||
|
||||
int32_t common_buffer_offset = 0;
|
||||
const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
|
||||
const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset +=
|
||||
cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
|
||||
const int32_t expanded_token_num = token_num * topk_num;
|
||||
const int32_t expand_token_id_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset +=
|
||||
cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
|
||||
const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset +=
|
||||
cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
|
||||
const int32_t input_quant_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset +=
|
||||
cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t));
|
||||
const int32_t input_scale_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float));
|
||||
const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += cpu_utils::round_up<64>(
|
||||
expanded_token_num * input_size_2 * sizeof(scalar_t));
|
||||
const int32_t w13_output_scale_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset +=
|
||||
cpu_utils::round_up<64>(expanded_token_num * sizeof(float));
|
||||
const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += cpu_utils::round_up<64>(
|
||||
expanded_token_num * output_size_2 * sizeof(float));
|
||||
|
||||
int32_t gemm_thread_buffer_offset = 0;
|
||||
const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset;
|
||||
gemm_thread_buffer_offset +=
|
||||
std::max(w13_input_buffer_size, w2_input_buffer_size);
|
||||
const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset;
|
||||
gemm_thread_buffer_offset += cpu_utils::round_up<64>(
|
||||
gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) *
|
||||
sizeof(int32_t));
|
||||
|
||||
const int32_t ws_output_buffer_offset = 0;
|
||||
const int32_t ws_thread_buffer_size =
|
||||
cpu_utils::round_up<64>(output_size_2 * sizeof(float));
|
||||
const int32_t thread_buffer_size =
|
||||
std::max(gemm_thread_buffer_offset, ws_thread_buffer_size);
|
||||
const int32_t buffer_size =
|
||||
common_buffer_offset + thread_buffer_size * thread_num;
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
|
||||
uint8_t* common_buffer_start =
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()
|
||||
->get_data<uint8_t>();
|
||||
uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset;
|
||||
|
||||
int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>(
|
||||
common_buffer_start + token_num_per_group_buffer_offset);
|
||||
int32_t* __restrict__ cu_token_num_per_group_buffer =
|
||||
reinterpret_cast<int32_t*>(common_buffer_start +
|
||||
cu_token_num_per_group_buffer_offset);
|
||||
int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>(
|
||||
common_buffer_start + expand_token_id_buffer_offset);
|
||||
int32_t* __restrict__ expand_token_id_index_buffer =
|
||||
reinterpret_cast<int32_t*>(common_buffer_start +
|
||||
expand_token_id_index_buffer_offset);
|
||||
int8_t* __restrict__ input_quant_buffer = reinterpret_cast<int8_t*>(
|
||||
common_buffer_start + input_quant_buffer_offset);
|
||||
float* __restrict__ input_scale_buffer =
|
||||
reinterpret_cast<float*>(common_buffer_start + input_scale_buffer_offset);
|
||||
|
||||
std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t));
|
||||
for (int32_t i = 0; i < expanded_token_num; ++i) {
|
||||
++token_num_per_group_buffer[topk_id[i]];
|
||||
}
|
||||
|
||||
int32_t token_num_sum = 0;
|
||||
cu_token_num_per_group_buffer[0] = 0;
|
||||
int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1;
|
||||
for (int32_t i = 0; i < expert_num; ++i) {
|
||||
token_index_buffer[i] = token_num_sum;
|
||||
token_num_sum += token_num_per_group_buffer[i];
|
||||
}
|
||||
|
||||
for (int32_t i = 0; i < token_num; ++i) {
|
||||
const int32_t* curr_topk_id = topk_id + i * topk_num;
|
||||
int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num;
|
||||
for (int32_t j = 0; j < topk_num; ++j) {
|
||||
const int32_t curr_expert_id = curr_topk_id[j];
|
||||
const int32_t curr_index = token_index_buffer[curr_expert_id]++;
|
||||
expand_token_id_buffer[curr_index] = i;
|
||||
curr_index_buffer[j] = curr_index;
|
||||
}
|
||||
}
|
||||
|
||||
// quantize inputs
|
||||
#pragma omp parallel for
|
||||
for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
gemm_t::quantize_row(input + token_idx * input_size_13,
|
||||
input_quant_buffer + token_idx * input_size_13,
|
||||
input_scale_buffer[token_idx], input_size_13);
|
||||
}
|
||||
|
||||
{
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
// w13 GEMM + act
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
const int32_t task_num_per_expert =
|
||||
(output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size;
|
||||
const int32_t task_num = task_num_per_expert * expert_num;
|
||||
uint8_t* __restrict__ thread_buffer =
|
||||
thread_buffer_start + thread_id * thread_buffer_size;
|
||||
int8_t* __restrict__ gemm_input_buffer =
|
||||
reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
|
||||
float* __restrict__ gemm_output_buffer =
|
||||
reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
|
||||
auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
|
||||
common_buffer_start + w13_gemm_output_buffer_offset);
|
||||
gemm_t gemm;
|
||||
|
||||
const int32_t w13_n_group_stride =
|
||||
gemm_t::WeightOCGroupSize * input_size_13;
|
||||
const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
|
||||
|
||||
for (;;) {
|
||||
const int32_t task_id = counter_ptr->acquire_counter();
|
||||
if (task_id >= task_num) {
|
||||
break;
|
||||
}
|
||||
const int32_t curr_expert_id = task_id / task_num_per_expert;
|
||||
const int32_t curr_output_group_id = task_id % task_num_per_expert;
|
||||
const int32_t curr_token_num =
|
||||
token_num_per_group_buffer[curr_expert_id];
|
||||
if (curr_token_num == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int32_t actual_n_tile_size =
|
||||
std::min(w13_n_tile_size,
|
||||
output_size_13 - curr_output_group_id * w13_n_tile_size);
|
||||
const int32_t* __restrict__ curr_expand_token_id_buffer =
|
||||
expand_token_id_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id];
|
||||
scalar_t* __restrict__ curr_w13_gemm_output_buffer =
|
||||
w13_gemm_output_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 +
|
||||
curr_output_group_id * w13_n_tile_size / 2;
|
||||
|
||||
const int8_t* w13_weight_ptr_0 = nullptr;
|
||||
const int8_t* w13_weight_ptr_1 = nullptr;
|
||||
const float* w13_scale_ptr_0 = nullptr;
|
||||
const float* w13_scale_ptr_1 = nullptr;
|
||||
scalar_t* w13_bias_ptr_0 = nullptr;
|
||||
scalar_t* w13_bias_ptr_1 = nullptr;
|
||||
if (act_type == FusedMOEAct::SwigluOAIAndMul) {
|
||||
const int32_t output_offset = curr_output_group_id * w13_n_tile_size;
|
||||
w13_weight_ptr_0 = w13 +
|
||||
curr_expert_id * input_size_13 * output_size_13 +
|
||||
output_offset * input_size_13;
|
||||
w13_weight_ptr_1 =
|
||||
w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13;
|
||||
w13_scale_ptr_0 =
|
||||
w13_scales + curr_expert_id * output_size_13 + output_offset;
|
||||
w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2;
|
||||
if (w13_bias != nullptr) {
|
||||
w13_bias_ptr_0 =
|
||||
w13_bias + curr_expert_id * output_size_13 + output_offset;
|
||||
w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2;
|
||||
}
|
||||
} else {
|
||||
const int32_t output_offset =
|
||||
curr_output_group_id * (w13_n_tile_size / 2);
|
||||
w13_weight_ptr_0 = w13 +
|
||||
curr_expert_id * input_size_13 * output_size_13 +
|
||||
output_offset * input_size_13;
|
||||
w13_weight_ptr_1 =
|
||||
w13_weight_ptr_0 + output_size_13 / 2 * input_size_13;
|
||||
w13_scale_ptr_0 =
|
||||
w13_scales + curr_expert_id * output_size_13 + output_offset;
|
||||
w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2;
|
||||
if (w13_bias != nullptr) {
|
||||
w13_bias_ptr_0 =
|
||||
w13_bias + curr_expert_id * output_size_13 + output_offset;
|
||||
w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2;
|
||||
}
|
||||
}
|
||||
|
||||
for (int32_t token_idx = 0; token_idx < curr_token_num;
|
||||
token_idx += gemm_m_tile_size) {
|
||||
const int32_t actual_token_num =
|
||||
std::min(gemm_m_tile_size, curr_token_num - token_idx);
|
||||
const int8_t* input_rows[gemm_m_tile_size];
|
||||
alignas(64) float input_scales[gemm_m_tile_size];
|
||||
// gather and pack
|
||||
for (int32_t i = 0; i < actual_token_num; ++i) {
|
||||
const int32_t curr_token_id = curr_expand_token_id_buffer[i];
|
||||
input_rows[i] = input_quant_buffer + curr_token_id * input_size_13;
|
||||
input_scales[i] = input_scale_buffer[curr_token_id];
|
||||
}
|
||||
gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
|
||||
actual_token_num, input_size_13);
|
||||
curr_expand_token_id_buffer += actual_token_num;
|
||||
|
||||
const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0;
|
||||
const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1;
|
||||
const float* w13_scale_ptr_0_iter = w13_scale_ptr_0;
|
||||
const float* w13_scale_ptr_1_iter = w13_scale_ptr_1;
|
||||
scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0;
|
||||
scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1;
|
||||
float* w13_output_buffer_0_iter = gemm_output_buffer;
|
||||
float* w13_output_buffer_1_iter =
|
||||
gemm_output_buffer + actual_n_tile_size / 2;
|
||||
|
||||
for (int32_t i = 0; i < actual_n_tile_size;
|
||||
i += min_w13_n_tile_size) {
|
||||
auto* output_0_int32 =
|
||||
reinterpret_cast<int32_t*>(w13_output_buffer_0_iter);
|
||||
gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32,
|
||||
actual_token_num, input_size_13, w13_n_group_stride,
|
||||
actual_n_tile_size);
|
||||
gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter,
|
||||
input_scales, w13_scale_ptr_0_iter,
|
||||
actual_token_num, gemm_n_tile_size,
|
||||
actual_n_tile_size);
|
||||
if (w13_bias != nullptr) {
|
||||
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
|
||||
w13_output_buffer_0_iter, w13_output_buffer_0_iter,
|
||||
w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size,
|
||||
actual_n_tile_size);
|
||||
w13_bias_ptr_0_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
auto* output_1_int32 =
|
||||
reinterpret_cast<int32_t*>(w13_output_buffer_1_iter);
|
||||
gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32,
|
||||
actual_token_num, input_size_13, w13_n_group_stride,
|
||||
actual_n_tile_size);
|
||||
gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter,
|
||||
input_scales, w13_scale_ptr_1_iter,
|
||||
actual_token_num, gemm_n_tile_size,
|
||||
actual_n_tile_size);
|
||||
if (w13_bias != nullptr) {
|
||||
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
|
||||
w13_output_buffer_1_iter, w13_output_buffer_1_iter,
|
||||
w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size,
|
||||
actual_n_tile_size);
|
||||
w13_bias_ptr_1_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
w13_weight_ptr_0_iter += w13_n_tile_stride;
|
||||
w13_weight_ptr_1_iter += w13_n_tile_stride;
|
||||
w13_scale_ptr_0_iter += gemm_n_tile_size;
|
||||
w13_scale_ptr_1_iter += gemm_n_tile_size;
|
||||
w13_output_buffer_0_iter += gemm_n_tile_size;
|
||||
w13_output_buffer_1_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
apply_gated_act(act_type, gemm_output_buffer,
|
||||
curr_w13_gemm_output_buffer, actual_token_num,
|
||||
actual_n_tile_size, actual_n_tile_size, input_size_2);
|
||||
curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
|
||||
common_buffer_start + w13_gemm_output_buffer_offset);
|
||||
float* __restrict__ w13_output_scale_buffer = reinterpret_cast<float*>(
|
||||
common_buffer_start + w13_output_scale_buffer_offset);
|
||||
|
||||
// quantize w2 inputs - in place
|
||||
#pragma omp parallel for
|
||||
for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) {
|
||||
scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2;
|
||||
int8_t* output_row = reinterpret_cast<int8_t*>(input_row);
|
||||
gemm_t::quantize_row(input_row, output_row,
|
||||
w13_output_scale_buffer[token_idx], input_size_2);
|
||||
}
|
||||
|
||||
{
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
// w2 gemm
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
const int32_t task_num_per_expert =
|
||||
(output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size;
|
||||
const int32_t task_num = task_num_per_expert * expert_num;
|
||||
uint8_t* __restrict__ thread_buffer =
|
||||
thread_buffer_start + thread_id * thread_buffer_size;
|
||||
int8_t* __restrict__ gemm_input_buffer =
|
||||
reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
|
||||
float* __restrict__ gemm_output_buffer =
|
||||
reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
|
||||
float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
|
||||
common_buffer_start + w2_gemm_output_buffer_offset);
|
||||
gemm_t gemm;
|
||||
|
||||
const int32_t w2_n_group_stride =
|
||||
gemm_t::WeightOCGroupSize * input_size_2;
|
||||
const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
|
||||
|
||||
for (;;) {
|
||||
const int32_t task_id = counter_ptr->acquire_counter();
|
||||
if (task_id >= task_num) {
|
||||
break;
|
||||
}
|
||||
const int32_t curr_expert_id = task_id / task_num_per_expert;
|
||||
const int32_t curr_output_group_id = task_id % task_num_per_expert;
|
||||
const int32_t curr_token_num =
|
||||
token_num_per_group_buffer[curr_expert_id];
|
||||
if (curr_token_num == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int32_t actual_n_tile_size =
|
||||
std::min(w2_n_tile_size,
|
||||
output_size_2 - curr_output_group_id * w2_n_tile_size);
|
||||
scalar_t* __restrict__ curr_w13_gemm_output_buffer =
|
||||
w13_gemm_output_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id] * input_size_2;
|
||||
float* __restrict__ curr_w13_output_scale_buffer =
|
||||
w13_output_scale_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id];
|
||||
float* __restrict__ curr_w2_gemm_output_buffer =
|
||||
w2_gemm_output_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size;
|
||||
const int8_t* __restrict__ w2_weight_ptr =
|
||||
w2 + curr_expert_id * output_size_2 * input_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size * input_size_2;
|
||||
const float* __restrict__ w2_scale_ptr =
|
||||
w2_scales + curr_expert_id * output_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size;
|
||||
scalar_t* w2_bias_ptr = nullptr;
|
||||
if (w2_bias != nullptr) {
|
||||
w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size;
|
||||
}
|
||||
|
||||
for (int32_t token_idx = 0; token_idx < curr_token_num;
|
||||
token_idx += gemm_m_tile_size) {
|
||||
const int32_t actual_token_num =
|
||||
std::min(gemm_m_tile_size, curr_token_num - token_idx);
|
||||
const int8_t* input_rows[gemm_m_tile_size];
|
||||
alignas(64) float input_scales[gemm_m_tile_size];
|
||||
for (int32_t i = 0; i < actual_token_num; ++i) {
|
||||
input_rows[i] = reinterpret_cast<const int8_t*>(
|
||||
curr_w13_gemm_output_buffer + i * input_size_2);
|
||||
input_scales[i] = curr_w13_output_scale_buffer[i];
|
||||
}
|
||||
gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
|
||||
actual_token_num, input_size_2);
|
||||
|
||||
const int8_t* w2_weight_ptr_iter = w2_weight_ptr;
|
||||
const float* w2_scale_ptr_iter = w2_scale_ptr;
|
||||
scalar_t* w2_bias_ptr_iter = w2_bias_ptr;
|
||||
float* curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer;
|
||||
for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
|
||||
auto* output_int32 = reinterpret_cast<int32_t*>(gemm_output_buffer);
|
||||
gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32,
|
||||
actual_token_num, input_size_2, w2_n_group_stride,
|
||||
gemm_n_tile_size);
|
||||
gemm_t::dequantize_tile(output_int32, gemm_output_buffer,
|
||||
input_scales, w2_scale_ptr_iter,
|
||||
actual_token_num, gemm_n_tile_size,
|
||||
gemm_n_tile_size);
|
||||
if (w2_bias != nullptr) {
|
||||
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
|
||||
gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter,
|
||||
actual_token_num, gemm_n_tile_size, gemm_n_tile_size);
|
||||
w2_bias_ptr_iter += gemm_n_tile_size;
|
||||
}
|
||||
for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) {
|
||||
std::memcpy(
|
||||
curr_w2_gemm_output_buffer_iter + m_idx * output_size_2,
|
||||
gemm_output_buffer + m_idx * gemm_n_tile_size,
|
||||
gemm_n_tile_size * sizeof(float));
|
||||
}
|
||||
|
||||
w2_weight_ptr_iter += w2_n_tile_stride;
|
||||
w2_scale_ptr_iter += gemm_n_tile_size;
|
||||
curr_w2_gemm_output_buffer_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
|
||||
curr_w13_output_scale_buffer += gemm_m_tile_size;
|
||||
curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
uint8_t* __restrict__ thread_buffer =
|
||||
thread_buffer_start + thread_id * thread_buffer_size;
|
||||
float* __restrict__ ws_output_buffer =
|
||||
reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset);
|
||||
float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
|
||||
common_buffer_start + w2_gemm_output_buffer_offset);
|
||||
|
||||
for (;;) {
|
||||
const int32_t token_id = counter_ptr->acquire_counter();
|
||||
if (token_id >= token_num) {
|
||||
break;
|
||||
}
|
||||
int32_t* __restrict__ curr_expand_token_id_index_buffer =
|
||||
expand_token_id_index_buffer + token_id * topk_num;
|
||||
const float* __restrict__ curr_weight =
|
||||
topk_weights + token_id * topk_num;
|
||||
const float first_weight = skip_weighted ? 1.0f : curr_weight[0];
|
||||
scalar_t* __restrict__ curr_output_buffer =
|
||||
output + token_id * output_size_2;
|
||||
|
||||
if (topk_num > 1) {
|
||||
int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
|
||||
float* w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
float* ws_output_buffer_iter = ws_output_buffer;
|
||||
vec_op::FP32Vec16 weight_vec(first_weight);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
(vec * weight_vec).save(ws_output_buffer_iter);
|
||||
w2_output_iter += 16;
|
||||
ws_output_buffer_iter += 16;
|
||||
}
|
||||
|
||||
for (int32_t idx = 1; idx < topk_num - 1; ++idx) {
|
||||
w2_output_idx = curr_expand_token_id_index_buffer[idx];
|
||||
w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
ws_output_buffer_iter = ws_output_buffer;
|
||||
weight_vec = vec_op::FP32Vec16(curr_weight[idx]);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
vec_op::FP32Vec16 sum(ws_output_buffer_iter);
|
||||
(sum + vec * weight_vec).save(ws_output_buffer_iter);
|
||||
w2_output_iter += 16;
|
||||
ws_output_buffer_iter += 16;
|
||||
}
|
||||
}
|
||||
|
||||
const int32_t last_idx = topk_num - 1;
|
||||
w2_output_idx = curr_expand_token_id_index_buffer[last_idx];
|
||||
w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
ws_output_buffer_iter = ws_output_buffer;
|
||||
scalar_t* curr_output_buffer_iter = curr_output_buffer;
|
||||
weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
vec_op::FP32Vec16 sum(ws_output_buffer_iter);
|
||||
scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter);
|
||||
w2_output_iter += 16;
|
||||
ws_output_buffer_iter += 16;
|
||||
curr_output_buffer_iter += 16;
|
||||
}
|
||||
} else {
|
||||
const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
|
||||
float* w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
scalar_t* curr_output_buffer_iter = curr_output_buffer;
|
||||
vec_op::FP32Vec16 weight_vec(first_weight);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter);
|
||||
w2_output_iter += 16;
|
||||
curr_output_buffer_iter += 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void prepack_moe_weight_int8(
|
||||
const torch::Tensor& weight, // [expert_num, output_size, input_size]
|
||||
torch::Tensor& packed_weight, const std::string& isa) {
|
||||
TORCH_CHECK(weight.is_contiguous());
|
||||
const int32_t expert_num = weight.size(0);
|
||||
const int32_t output_size = weight.size(1);
|
||||
const int32_t input_size = weight.size(2);
|
||||
const int64_t expert_stride = weight.stride(0);
|
||||
const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
|
||||
TORCH_CHECK_EQ(output_size % 32, 0);
|
||||
|
||||
CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() {
|
||||
TORCH_CHECK_EQ(input_size % gemm_t::K, 0);
|
||||
prepack_moe_weight_int8_impl<gemm_t>(
|
||||
weight.data_ptr<int8_t>(), packed_weight.data_ptr<int8_t>(), expert_num,
|
||||
output_size, input_size, expert_stride);
|
||||
});
|
||||
}
|
||||
|
||||
void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
|
||||
const torch::Tensor& w13, const torch::Tensor& w2,
|
||||
const torch::Tensor& w13_scale,
|
||||
const torch::Tensor& w2_scale,
|
||||
const std::optional<torch::Tensor>& w13_bias,
|
||||
const std::optional<torch::Tensor>& w2_bias,
|
||||
const torch::Tensor& topk_weights,
|
||||
const torch::Tensor& topk_id, const bool skip_weighted,
|
||||
const std::string& act, const std::string& isa) {
|
||||
const int32_t token_num = input.size(0);
|
||||
const int32_t input_size_13 = input.size(1);
|
||||
const int64_t input_stride = input.stride(0);
|
||||
TORCH_CHECK_EQ(input_stride, input_size_13);
|
||||
const int32_t expert_num = w13.size(0);
|
||||
const int32_t output_size_13 = w13.size(1);
|
||||
const int32_t input_size_2 = w2.size(2);
|
||||
const int32_t output_size_2 = w2.size(1);
|
||||
const int32_t topk_num = topk_id.size(1);
|
||||
const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
|
||||
const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
|
||||
TORCH_CHECK(!skip_weighted || topk_num == 1,
|
||||
"skip_weighted is only supported for topk=1 on CPU");
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "cpu_fused_moe_int8", [&]() {
|
||||
CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() {
|
||||
fused_moe_int8_impl<scalar_t, gemm_t>(
|
||||
output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
|
||||
w13.data_ptr<int8_t>(), w2.data_ptr<int8_t>(),
|
||||
w13_scale.data_ptr<float>(), w2_scale.data_ptr<float>(),
|
||||
w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr,
|
||||
w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
|
||||
topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(),
|
||||
act_type, token_num, expert_num, topk_num, input_size_13,
|
||||
output_size_13, input_size_2, output_size_2, skip_weighted);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
|
||||
explicit FP32Vec4(__vector float data) : reg(data) {}
|
||||
|
||||
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
|
||||
FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
|
||||
};
|
||||
|
||||
struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
@@ -316,7 +316,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
|
||||
explicit FP32Vec8(f32x4x2_t data) : reg(data) {}
|
||||
|
||||
explicit FP32Vec8(const FP32Vec8& data) {
|
||||
FP32Vec8(const FP32Vec8& data) {
|
||||
reg.val[0] = data.reg.val[0];
|
||||
reg.val[1] = data.reg.val[1];
|
||||
}
|
||||
@@ -593,7 +593,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
|
||||
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
|
||||
|
||||
explicit FP32Vec16(const FP32Vec16& data) {
|
||||
FP32Vec16(const FP32Vec16& data) {
|
||||
reg.val[0] = data.reg.val[0];
|
||||
reg.val[1] = data.reg.val[1];
|
||||
reg.val[2] = data.reg.val[2];
|
||||
@@ -747,6 +747,15 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
|
||||
}
|
||||
|
||||
FP32Vec16 exp() const {
|
||||
FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
|
||||
FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
|
||||
auto lo_e = lo.exp();
|
||||
auto hi_e = hi.exp();
|
||||
return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
|
||||
hi_e.reg.val[0], hi_e.reg.val[1]});
|
||||
}
|
||||
|
||||
float reduce_max() {
|
||||
__vector float max01 = vec_max(reg.val[0], reg.val[1]);
|
||||
__vector float max23 = vec_max(reg.val[2], reg.val[3]);
|
||||
|
||||
@@ -31,6 +31,9 @@ class MicroGemm {
|
||||
}
|
||||
};
|
||||
|
||||
template <cpu_utils::ISA isa, typename scalar_t>
|
||||
class MicroGemmINT8;
|
||||
|
||||
template <int32_t n_size, typename scalar_t>
|
||||
FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
|
||||
scalar_t* __restrict__ d_ptr,
|
||||
|
||||
@@ -0,0 +1,424 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#ifndef CPU_MICRO_GEMM_INT8_NEON_HPP
|
||||
#define CPU_MICRO_GEMM_INT8_NEON_HPP
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
#include <arm_bf16.h>
|
||||
#include <arm_neon.h>
|
||||
#include <c10/util/BFloat16.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/Half.h>
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
|
||||
namespace neon_smmla {
|
||||
|
||||
constexpr int32_t K = 8;
|
||||
constexpr int32_t Cols = 2;
|
||||
constexpr int32_t TileSize = K * Cols;
|
||||
|
||||
FORCE_INLINE float32x4x2_t load_as_f32(const float* input) {
|
||||
float32x4x2_t result;
|
||||
result.val[0] = vld1q_f32(input);
|
||||
result.val[1] = vld1q_f32(input + 4);
|
||||
return result;
|
||||
}
|
||||
|
||||
FORCE_INLINE float32x4x2_t load_as_f32(const c10::Half* input) {
|
||||
const auto input_vec = vld1q_f16(reinterpret_cast<const float16_t*>(input));
|
||||
float32x4x2_t result;
|
||||
result.val[0] = vcvt_f32_f16(vget_low_f16(input_vec));
|
||||
result.val[1] = vcvt_f32_f16(vget_high_f16(input_vec));
|
||||
return result;
|
||||
}
|
||||
|
||||
FORCE_INLINE float32x4x2_t load_as_f32(const c10::BFloat16* input) {
|
||||
const auto input_vec = vld1q_bf16(reinterpret_cast<const bfloat16_t*>(input));
|
||||
float32x4x2_t result;
|
||||
result.val[0] = vcvt_f32_bf16(vget_low_bf16(input_vec));
|
||||
result.val[1] = vcvt_f32_bf16(vget_high_bf16(input_vec));
|
||||
return result;
|
||||
}
|
||||
|
||||
FORCE_INLINE void store_acc_rowpair(const int32x4_t acc01,
|
||||
const int32x4_t acc23,
|
||||
const int32x4_t acc45,
|
||||
const int32x4_t acc67,
|
||||
int32_t* __restrict__ c_ptr,
|
||||
const int64_t ldc, const int32_t m_rows) {
|
||||
if (m_rows == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
vst1q_s32(c_ptr, vcombine_s32(vget_low_s32(acc01), vget_low_s32(acc23)));
|
||||
vst1q_s32(c_ptr + 4, vcombine_s32(vget_low_s32(acc45), vget_low_s32(acc67)));
|
||||
|
||||
if (m_rows == 2) {
|
||||
vst1q_s32(c_ptr + ldc,
|
||||
vcombine_s32(vget_high_s32(acc01), vget_high_s32(acc23)));
|
||||
vst1q_s32(c_ptr + ldc + 4,
|
||||
vcombine_s32(vget_high_s32(acc45), vget_high_s32(acc67)));
|
||||
}
|
||||
}
|
||||
|
||||
FORCE_INLINE void gemm_micro_smmla_8x8_packed_a(
|
||||
const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
|
||||
int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
|
||||
const int64_t ldc) {
|
||||
const int32x4_t zero = vdupq_n_s32(0);
|
||||
int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
|
||||
int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
|
||||
int32x4_t acc4501 = zero, acc4523 = zero, acc4545 = zero, acc4567 = zero;
|
||||
int32x4_t acc6701 = zero, acc6723 = zero, acc6745 = zero, acc6767 = zero;
|
||||
|
||||
const int8_t* __restrict__ a_tile = a_packed;
|
||||
const int8_t* __restrict__ b_tile = b_packed;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
|
||||
const int8x16_t a_tile01 = vld1q_s8(a_tile);
|
||||
const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
|
||||
const int8x16_t a_tile45 = vld1q_s8(a_tile + 2 * TileSize);
|
||||
const int8x16_t a_tile67 = vld1q_s8(a_tile + 3 * TileSize);
|
||||
|
||||
const int8x16_t b_tile01 = vld1q_s8(b_tile);
|
||||
const int8x16_t b_tile23 = vld1q_s8(b_tile + TileSize);
|
||||
const int8x16_t b_tile45 = vld1q_s8(b_tile + 2 * TileSize);
|
||||
const int8x16_t b_tile67 = vld1q_s8(b_tile + 3 * TileSize);
|
||||
|
||||
acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
|
||||
acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
|
||||
acc4501 = vmmlaq_s32(acc4501, a_tile45, b_tile01);
|
||||
acc6701 = vmmlaq_s32(acc6701, a_tile67, b_tile01);
|
||||
|
||||
acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
|
||||
acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
|
||||
acc4523 = vmmlaq_s32(acc4523, a_tile45, b_tile23);
|
||||
acc6723 = vmmlaq_s32(acc6723, a_tile67, b_tile23);
|
||||
|
||||
acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
|
||||
acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
|
||||
acc4545 = vmmlaq_s32(acc4545, a_tile45, b_tile45);
|
||||
acc6745 = vmmlaq_s32(acc6745, a_tile67, b_tile45);
|
||||
|
||||
acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
|
||||
acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
|
||||
acc4567 = vmmlaq_s32(acc4567, a_tile45, b_tile67);
|
||||
acc6767 = vmmlaq_s32(acc6767, a_tile67, b_tile67);
|
||||
|
||||
a_tile += 4 * TileSize;
|
||||
b_tile += 4 * TileSize;
|
||||
}
|
||||
|
||||
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
|
||||
std::min(2, m));
|
||||
store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 2)));
|
||||
store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 4)));
|
||||
store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 6)));
|
||||
}
|
||||
|
||||
FORCE_INLINE void gemm_micro_smmla_4x16_packed_a(
|
||||
const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
|
||||
int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
|
||||
const int64_t b_n_group_stride, const int64_t ldc) {
|
||||
const int32_t m_rows_01 = std::min(2, m);
|
||||
const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
|
||||
const int32x4_t zero = vdupq_n_s32(0);
|
||||
|
||||
int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
|
||||
int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
|
||||
int32x4_t acc0189 = zero, acc011011 = zero, acc011213 = zero,
|
||||
acc011415 = zero;
|
||||
int32x4_t acc2389 = zero, acc231011 = zero, acc231213 = zero,
|
||||
acc231415 = zero;
|
||||
|
||||
const int8_t* __restrict__ a_tile = a_packed;
|
||||
// note: b packs 8 panels contiguously, so we need 2 b_tile ptrs
|
||||
// for the 4x16 microkernel
|
||||
const int8_t* __restrict__ b_tile0 = b_packed;
|
||||
const int8_t* __restrict__ b_tile1 = b_packed + b_n_group_stride;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
|
||||
const int8x16_t a_tile01 = vld1q_s8(a_tile);
|
||||
const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
|
||||
const int8x16_t b_tile01 = vld1q_s8(b_tile0);
|
||||
const int8x16_t b_tile23 = vld1q_s8(b_tile0 + TileSize);
|
||||
const int8x16_t b_tile45 = vld1q_s8(b_tile0 + 2 * TileSize);
|
||||
const int8x16_t b_tile67 = vld1q_s8(b_tile0 + 3 * TileSize);
|
||||
const int8x16_t b_tile89 = vld1q_s8(b_tile1);
|
||||
const int8x16_t b_tile1011 = vld1q_s8(b_tile1 + TileSize);
|
||||
const int8x16_t b_tile1213 = vld1q_s8(b_tile1 + 2 * TileSize);
|
||||
const int8x16_t b_tile1415 = vld1q_s8(b_tile1 + 3 * TileSize);
|
||||
|
||||
acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
|
||||
acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
|
||||
acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
|
||||
acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
|
||||
|
||||
acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
|
||||
acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
|
||||
acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
|
||||
acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
|
||||
|
||||
acc0189 = vmmlaq_s32(acc0189, a_tile01, b_tile89);
|
||||
acc2389 = vmmlaq_s32(acc2389, a_tile23, b_tile89);
|
||||
acc011011 = vmmlaq_s32(acc011011, a_tile01, b_tile1011);
|
||||
acc231011 = vmmlaq_s32(acc231011, a_tile23, b_tile1011);
|
||||
|
||||
acc011213 = vmmlaq_s32(acc011213, a_tile01, b_tile1213);
|
||||
acc231213 = vmmlaq_s32(acc231213, a_tile23, b_tile1213);
|
||||
acc011415 = vmmlaq_s32(acc011415, a_tile01, b_tile1415);
|
||||
acc231415 = vmmlaq_s32(acc231415, a_tile23, b_tile1415);
|
||||
|
||||
a_tile += 2 * TileSize;
|
||||
b_tile0 += 4 * TileSize;
|
||||
b_tile1 += 4 * TileSize;
|
||||
}
|
||||
|
||||
// rows 0-1, columns 0-7
|
||||
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
|
||||
// rows 0-1, columns 8-15
|
||||
store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
|
||||
m_rows_01);
|
||||
// rows 2-3, columns 0-7
|
||||
store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
|
||||
m_rows_23);
|
||||
// rows 2-3, columns 8-15
|
||||
store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
|
||||
c_ptr + 2 * ldc + 8, ldc, m_rows_23);
|
||||
}
|
||||
|
||||
} // namespace neon_smmla
|
||||
|
||||
template <typename scalar_t>
|
||||
class MicroGemmINT8<cpu_utils::ISA::NEON, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t K = neon_smmla::K;
|
||||
static constexpr int32_t Mr = 8;
|
||||
static constexpr int32_t Nr = 8;
|
||||
static constexpr int32_t NrGemv = 16;
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = Nr;
|
||||
static_assert(MaxMSize % Mr == 0);
|
||||
|
||||
static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output,
|
||||
float& scale, const int32_t size) {
|
||||
TORCH_CHECK_EQ(size % K, 0);
|
||||
float32x4_t max_vec = vdupq_n_f32(0.0f);
|
||||
|
||||
for (int32_t i = 0; i < size; i += K) {
|
||||
const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
|
||||
max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[0]));
|
||||
max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[1]));
|
||||
}
|
||||
|
||||
const float abs_max = std::max(vmaxvq_f32(max_vec), 1.0e-7f);
|
||||
scale = abs_max / 127.0f;
|
||||
const float32x4_t inv_scale_vec = vdupq_n_f32(127.0f / abs_max);
|
||||
|
||||
for (int32_t i = 0; i < size; i += K) {
|
||||
const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
|
||||
const int32x4_t output_low =
|
||||
vcvtnq_s32_f32(vmulq_f32(input_vec.val[0], inv_scale_vec));
|
||||
const int32x4_t output_high =
|
||||
vcvtnq_s32_f32(vmulq_f32(input_vec.val[1], inv_scale_vec));
|
||||
const int16x8_t output_s16 =
|
||||
vcombine_s16(vqmovn_s32(output_low), vqmovn_s32(output_high));
|
||||
vst1_s8(output + i, vqmovn_s16(output_s16));
|
||||
}
|
||||
}
|
||||
|
||||
// with current code, fusing this into the gemm micro kernel didn't move the
|
||||
// needle
|
||||
static FORCE_INLINE void dequantize_tile(
|
||||
int32_t* input, float* output, const float* __restrict__ input_scales,
|
||||
const float* __restrict__ weight_scales, const int32_t m, const int32_t n,
|
||||
const int32_t stride) {
|
||||
TORCH_CHECK_EQ(n % 4, 0);
|
||||
for (int32_t m_idx = 0; m_idx < m; ++m_idx) {
|
||||
const float32x4_t input_scale_vec = vdupq_n_f32(input_scales[m_idx]);
|
||||
for (int32_t n_idx = 0; n_idx < n; n_idx += 4) {
|
||||
const int32x4_t input_vec = vld1q_s32(input + m_idx * stride + n_idx);
|
||||
const float32x4_t weight_scale_vec = vld1q_f32(weight_scales + n_idx);
|
||||
const float32x4_t output_vec =
|
||||
vmulq_f32(vcvtq_f32_s32(input_vec),
|
||||
vmulq_f32(input_scale_vec, weight_scale_vec));
|
||||
vst1q_f32(output + m_idx * stride + n_idx, output_vec);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// physical layout [
|
||||
// M / (8 or 4); Mr is 8 or 4
|
||||
// K / 8; K for smmla is 8
|
||||
// 4, ; 4 row-pairs for each 8 rows
|
||||
// 2, ; row-pair is 2 rows
|
||||
// 4 ; 4 elements per row
|
||||
// ]
|
||||
static void pack_input_from_rows(const int8_t* const* __restrict__ rows,
|
||||
int8_t* __restrict__ a_packed,
|
||||
const int32_t m, const int32_t k) {
|
||||
TORCH_CHECK(m > 0 && m <= MaxMSize);
|
||||
TORCH_CHECK(k % K == 0);
|
||||
const int8x8_t zero = vdup_n_s8(0);
|
||||
|
||||
for (int32_t row_base = 0; row_base < m; row_base += Mr) {
|
||||
const int32_t panel_m = std::min(Mr, m - row_base);
|
||||
const int8_t* const* panel_rows = rows + row_base;
|
||||
int8_t* __restrict__ out = a_packed + row_base * k;
|
||||
|
||||
// fast path for full 8-row panels (fast path for 4-row panels didn't move
|
||||
// the needle)
|
||||
if (panel_m == Mr) {
|
||||
const int8_t* __restrict__ row0 = panel_rows[0];
|
||||
const int8_t* __restrict__ row1 = panel_rows[1];
|
||||
const int8_t* __restrict__ row2 = panel_rows[2];
|
||||
const int8_t* __restrict__ row3 = panel_rows[3];
|
||||
const int8_t* __restrict__ row4 = panel_rows[4];
|
||||
const int8_t* __restrict__ row5 = panel_rows[5];
|
||||
const int8_t* __restrict__ row6 = panel_rows[6];
|
||||
const int8_t* __restrict__ row7 = panel_rows[7];
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 2 * K <= k; k_idx += 2 * K) {
|
||||
int8_t* __restrict__ block0 = out;
|
||||
int8_t* __restrict__ block1 = out + 4 * neon_smmla::TileSize;
|
||||
|
||||
int8x16_t a0 = vld1q_s8(row0 + k_idx);
|
||||
int8x16_t a1 = vld1q_s8(row1 + k_idx);
|
||||
vst1q_s8(block0, vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
|
||||
vst1q_s8(block1, vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
|
||||
|
||||
a0 = vld1q_s8(row2 + k_idx);
|
||||
a1 = vld1q_s8(row3 + k_idx);
|
||||
vst1q_s8(block0 + neon_smmla::TileSize,
|
||||
vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
|
||||
vst1q_s8(block1 + neon_smmla::TileSize,
|
||||
vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
|
||||
|
||||
a0 = vld1q_s8(row4 + k_idx);
|
||||
a1 = vld1q_s8(row5 + k_idx);
|
||||
vst1q_s8(block0 + 2 * neon_smmla::TileSize,
|
||||
vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
|
||||
vst1q_s8(block1 + 2 * neon_smmla::TileSize,
|
||||
vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
|
||||
|
||||
a0 = vld1q_s8(row6 + k_idx);
|
||||
a1 = vld1q_s8(row7 + k_idx);
|
||||
vst1q_s8(block0 + 3 * neon_smmla::TileSize,
|
||||
vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
|
||||
vst1q_s8(block1 + 3 * neon_smmla::TileSize,
|
||||
vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
|
||||
|
||||
out += 8 * neon_smmla::TileSize;
|
||||
}
|
||||
|
||||
for (; k_idx < k; k_idx += K) {
|
||||
int8x8_t a0 = vld1_s8(row0 + k_idx);
|
||||
int8x8_t a1 = vld1_s8(row1 + k_idx);
|
||||
vst1q_s8(out, vcombine_s8(a0, a1));
|
||||
|
||||
a0 = vld1_s8(row2 + k_idx);
|
||||
a1 = vld1_s8(row3 + k_idx);
|
||||
vst1q_s8(out + neon_smmla::TileSize, vcombine_s8(a0, a1));
|
||||
|
||||
a0 = vld1_s8(row4 + k_idx);
|
||||
a1 = vld1_s8(row5 + k_idx);
|
||||
vst1q_s8(out + 2 * neon_smmla::TileSize, vcombine_s8(a0, a1));
|
||||
|
||||
a0 = vld1_s8(row6 + k_idx);
|
||||
a1 = vld1_s8(row7 + k_idx);
|
||||
vst1q_s8(out + 3 * neon_smmla::TileSize, vcombine_s8(a0, a1));
|
||||
|
||||
out += 4 * neon_smmla::TileSize;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const int32_t row_pairs = (panel_m <= 4) ? 2 : Mr / 2;
|
||||
for (int32_t k_idx = 0; k_idx < k; k_idx += K) {
|
||||
for (int32_t pair_idx = 0; pair_idx < row_pairs; ++pair_idx) {
|
||||
const int32_t row_idx = pair_idx * 2;
|
||||
const int8x8_t row0 =
|
||||
(row_idx < panel_m) ? vld1_s8(panel_rows[row_idx] + k_idx) : zero;
|
||||
const int8x8_t row1 = (row_idx + 1 < panel_m)
|
||||
? vld1_s8(panel_rows[row_idx + 1] + k_idx)
|
||||
: zero;
|
||||
vst1q_s8(out, vcombine_s8(row0, row1));
|
||||
out += neon_smmla::TileSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// physical layout [
|
||||
// N / 8; Nr is 8
|
||||
// K / 8; K for smmla is 8
|
||||
// 4, ; 4 col-pairs for each 8 cols
|
||||
// 2, ; col-pair is 2 cols
|
||||
// 4 ; 4 elements per col
|
||||
// ]
|
||||
static void pack_weight(const int8_t* __restrict__ weight,
|
||||
int8_t* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
TORCH_CHECK(output_size % NSize == 0);
|
||||
TORCH_CHECK(input_size % K == 0);
|
||||
|
||||
for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
|
||||
int8_t* __restrict__ dst = packed_weight + o_idx * input_size;
|
||||
for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) {
|
||||
for (int32_t pair_idx = 0; pair_idx < Nr;
|
||||
pair_idx += neon_smmla::Cols) {
|
||||
const int8_t* __restrict__ row0 =
|
||||
weight + (o_idx + pair_idx) * input_size + k_idx;
|
||||
const int8_t* __restrict__ row1 = row0 + input_size;
|
||||
vst1q_s8(dst, vcombine_s8(vld1_s8(row0), vld1_s8(row1)));
|
||||
dst += neon_smmla::TileSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void gemm(const int8_t* __restrict__ a_packed,
|
||||
const int8_t* __restrict__ b_packed, int32_t* __restrict__ c,
|
||||
const int32_t m, const int32_t k, const int64_t b_n_group_stride,
|
||||
const int64_t ldc) const {
|
||||
TORCH_CHECK(m > 0 && m <= MaxMSize);
|
||||
TORCH_CHECK(k % K == 0);
|
||||
|
||||
for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
|
||||
const int8_t* __restrict__ b_panel = b_packed + n_idx * k;
|
||||
|
||||
for (int32_t row_base = 0; row_base < m; row_base += Mr) {
|
||||
const int32_t panel_m = std::min(Mr, m - row_base);
|
||||
const int8_t* __restrict__ a_panel = a_packed + row_base * k;
|
||||
int32_t* __restrict__ c_panel = c + row_base * ldc + n_idx;
|
||||
|
||||
if (panel_m <= 4) {
|
||||
neon_smmla::gemm_micro_smmla_4x16_packed_a(
|
||||
a_panel, b_panel, c_panel, panel_m, k, b_n_group_stride, ldc);
|
||||
} else {
|
||||
neon_smmla::gemm_micro_smmla_8x8_packed_a(a_panel, b_panel, c_panel,
|
||||
panel_m, k, ldc);
|
||||
neon_smmla::gemm_micro_smmla_8x8_packed_a(
|
||||
a_panel, b_panel + b_n_group_stride, c_panel + Nr, panel_m, k,
|
||||
ldc);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif
|
||||
@@ -1,3 +1,6 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#ifndef CPU_MICRO_GEMM_NEON_HPP
|
||||
#define CPU_MICRO_GEMM_NEON_HPP
|
||||
|
||||
@@ -16,9 +19,6 @@ namespace {
|
||||
constexpr int32_t K = 4;
|
||||
constexpr int32_t Cols = 2;
|
||||
constexpr int32_t TileSize = K * Cols;
|
||||
constexpr int32_t Mr = 8;
|
||||
constexpr int32_t Nr = 8;
|
||||
constexpr int32_t Nr_gemv = 16;
|
||||
|
||||
// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1]
|
||||
FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) {
|
||||
@@ -132,7 +132,7 @@ FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
|
||||
acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);
|
||||
|
||||
a_tile += 4 * TileSize;
|
||||
b_tile += Nr * K;
|
||||
b_tile += 4 * TileSize;
|
||||
}
|
||||
|
||||
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
|
||||
@@ -205,8 +205,8 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
|
||||
acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415);
|
||||
|
||||
a_tile += 2 * TileSize;
|
||||
b_tile0 += Nr * K;
|
||||
b_tile1 += Nr * K;
|
||||
b_tile0 += 4 * TileSize;
|
||||
b_tile1 += 4 * TileSize;
|
||||
}
|
||||
|
||||
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
|
||||
@@ -223,6 +223,9 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t Mr = 8;
|
||||
static constexpr int32_t Nr = 8;
|
||||
static constexpr int32_t NrGemv = 16;
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = Nr;
|
||||
@@ -246,6 +249,9 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
|
||||
public:
|
||||
using scalar_t = c10::BFloat16;
|
||||
|
||||
static constexpr int32_t Mr = 8;
|
||||
static constexpr int32_t Nr = 8;
|
||||
static constexpr int32_t NrGemv = 16;
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = Nr;
|
||||
@@ -253,7 +259,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
|
||||
|
||||
public:
|
||||
// physical layout [
|
||||
// M / 8; Mr is 8
|
||||
// M / (8 or 4); Mr is 8 or 4
|
||||
// K / 4; K for bfmmla is 4
|
||||
// 4, ; 4 row-pairs for each 8 rows
|
||||
// 2, ; row-pair is 2 rows
|
||||
@@ -439,7 +445,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
|
||||
(void)lda; // A is packed, so lda is not needed
|
||||
TORCH_CHECK_EQ(k % K, 0);
|
||||
|
||||
for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
|
||||
for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
|
||||
const bfloat16_t* __restrict__ b_panel =
|
||||
reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
|
||||
|
||||
|
||||
+173
-12
@@ -451,6 +451,90 @@ void causal_conv1d_update_kernel_impl(
|
||||
});
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void causal_conv1d_update_multi_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
scalar_t* __restrict__ conv_states,
|
||||
const scalar_t* __restrict__ weight,
|
||||
const scalar_t* __restrict__ bias,
|
||||
const int32_t* __restrict__ num_accepted_tokens,
|
||||
const int32_t* __restrict__ conv_indices,
|
||||
bool silu_activation,
|
||||
int64_t batch,
|
||||
int64_t dim,
|
||||
int64_t seqlen,
|
||||
int64_t width,
|
||||
int64_t state_len,
|
||||
int64_t conv_state_slot_stride) {
|
||||
constexpr int64_t BLOCK_N = block_size_n() * 2;
|
||||
const int64_t NB = div_up(dim, BLOCK_N);
|
||||
|
||||
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
|
||||
at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t bs{0}, nb{0};
|
||||
data_index_init(begin, bs, batch, nb, NB);
|
||||
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
const int64_t nb_start = nb * BLOCK_N;
|
||||
const int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
|
||||
const int32_t conv_state_index = conv_indices[bs];
|
||||
const int32_t history_offset = num_accepted_tokens[bs] - 1;
|
||||
|
||||
switch (width << 4 | nb_size >> 4) {
|
||||
case 0x42:
|
||||
tinygemm_kernel<scalar_t, 4, 32, has_bias, has_silu>::apply(
|
||||
input + bs * seqlen * dim + nb_start,
|
||||
weight + nb_start * width,
|
||||
out + bs * seqlen * dim + nb_start,
|
||||
has_bias ? bias + nb_start : nullptr,
|
||||
conv_states + conv_state_index * conv_state_slot_stride +
|
||||
history_offset * dim + nb_start,
|
||||
true,
|
||||
seqlen,
|
||||
dim,
|
||||
true);
|
||||
break;
|
||||
case 0x44:
|
||||
tinygemm_kernel<scalar_t, 4, 64, has_bias, has_silu>::apply(
|
||||
input + bs * seqlen * dim + nb_start,
|
||||
weight + nb_start * width,
|
||||
out + bs * seqlen * dim + nb_start,
|
||||
has_bias ? bias + nb_start : nullptr,
|
||||
conv_states + conv_state_index * conv_state_slot_stride +
|
||||
history_offset * dim + nb_start,
|
||||
true,
|
||||
seqlen,
|
||||
dim,
|
||||
true);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
|
||||
}
|
||||
|
||||
data_index_step(bs, batch, nb, NB);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t bs = begin; bs < end; ++bs) {
|
||||
const int32_t conv_state_index = conv_indices[bs];
|
||||
const int32_t num_accepted = num_accepted_tokens[bs];
|
||||
scalar_t* state = conv_states + conv_state_index * conv_state_slot_stride;
|
||||
|
||||
std::memmove(
|
||||
state,
|
||||
state + num_accepted * dim,
|
||||
(state_len - seqlen) * dim * sizeof(scalar_t));
|
||||
std::memcpy(
|
||||
state + (state_len - seqlen) * dim,
|
||||
input + bs * seqlen * dim,
|
||||
seqlen * dim * sizeof(scalar_t));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// from [dim, width] or [N, K]
|
||||
@@ -545,7 +629,7 @@ at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t n
|
||||
// query_start_loc: (batch + 1) int32
|
||||
// cache_indices: (batch) int32
|
||||
// has_initial_state: (batch) bool
|
||||
// conv_states: (..., dim, width - 1) itype
|
||||
// conv_states: (..., dim, state_len) itype, where state_len >= width - 1
|
||||
// activation: either None or "silu" or "swish"
|
||||
// pad_slot_id: int
|
||||
//
|
||||
@@ -586,11 +670,14 @@ at::Tensor causal_conv1d_fwd_cpu(
|
||||
CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
|
||||
CHECK_GE(padded_batch, batch);
|
||||
CHECK_EQ(conv_states_val.size(1), dim);
|
||||
CHECK_EQ(conv_states_val.size(2), width - 1);
|
||||
const int64_t state_len = conv_states_val.size(2);
|
||||
CHECK_GE(state_len, width - 1);
|
||||
|
||||
// adjust `conv_states` to be contiguous on `dim`
|
||||
// should happen only once
|
||||
if (conv_states_val.stride(-2) != 1) {
|
||||
TORCH_CHECK(state_len == width - 1,
|
||||
"causal_conv1d_fwd_cpu: wide conv_states must be contiguous on dim.");
|
||||
auto conv_states_copy = conv_states_val.clone();
|
||||
conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
|
||||
conv_states_val.copy_(conv_states_copy);
|
||||
@@ -651,14 +738,14 @@ at::Tensor causal_conv1d_fwd_cpu(
|
||||
|
||||
// API aligned with GPUs
|
||||
//
|
||||
// x: (batch, dim) or (batch, dim, seqlen)
|
||||
// x: (batch, dim) or (batch, seqlen, dim)
|
||||
// conv_state: (..., dim, state_len), where state_len >= width - 1
|
||||
// weight: (dim, width)
|
||||
// bias: (dim,)
|
||||
// cache_seqlens: (batch,), dtype int32.
|
||||
// num_accepted_tokens: (batch,), dtype int32.
|
||||
// conv_state_indices: (batch,), dtype int32
|
||||
// pad_slot_id: int
|
||||
// out: (batch, dim) or (batch, dim, seqlen)
|
||||
// out: (batch, dim) or (batch, seqlen, dim)
|
||||
//
|
||||
at::Tensor causal_conv1d_update_cpu(
|
||||
const at::Tensor& x,
|
||||
@@ -666,7 +753,7 @@ at::Tensor causal_conv1d_update_cpu(
|
||||
const at::Tensor& weight,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
bool silu_activation,
|
||||
const std::optional<at::Tensor>& cache_seqlens,
|
||||
const std::optional<at::Tensor>& num_accepted_tokens,
|
||||
const std::optional<at::Tensor>& conv_state_indices,
|
||||
int64_t pad_slot_id,
|
||||
bool is_vnni) {
|
||||
@@ -674,13 +761,13 @@ at::Tensor causal_conv1d_update_cpu(
|
||||
CHECK_CONTIGUOUS(weight);
|
||||
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
|
||||
|
||||
// TODO: add multi-token prediction support
|
||||
TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
|
||||
TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");
|
||||
TORCH_CHECK(
|
||||
x.dim() == 2 || x.dim() == 3,
|
||||
"causal_conv1d_update_cpu: expect x to be 2D or 3D tensor.");
|
||||
|
||||
int64_t batch = x.size(0);
|
||||
int64_t dim = x.size(1);
|
||||
int64_t seqlen = 1;
|
||||
int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2);
|
||||
int64_t seqlen = x.dim() == 2 ? 1 : x.size(1);
|
||||
int64_t width = weight.size(-1);
|
||||
|
||||
const auto scalar_type = x.scalar_type();
|
||||
@@ -690,10 +777,84 @@ at::Tensor causal_conv1d_update_cpu(
|
||||
|
||||
CHECK_EQ(conv_states.scalar_type(), scalar_type);
|
||||
CHECK_EQ(conv_states.size(1), dim);
|
||||
CHECK_EQ(conv_states.size(2), width - 1);
|
||||
const int64_t state_len = conv_states.size(2);
|
||||
CHECK_GE(state_len, width - 1);
|
||||
|
||||
if (x.dim() == 3) {
|
||||
TORCH_CHECK(
|
||||
num_accepted_tokens.has_value(),
|
||||
"causal_conv1d_update_cpu: num_accepted_tokens is required for 3D x.");
|
||||
TORCH_CHECK(
|
||||
conv_state_indices.has_value(),
|
||||
"causal_conv1d_update_cpu: conv_state_indices is required for 3D x.");
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(num_accepted_tokens, batch, at::kInt);
|
||||
TORCH_CHECK(
|
||||
width == 4,
|
||||
"causal_conv1d_update_cpu: support only width of 4 for 3D x.");
|
||||
TORCH_CHECK(
|
||||
seqlen > 0,
|
||||
"causal_conv1d_update_cpu: expect non-empty sequence for 3D x.");
|
||||
TORCH_CHECK(
|
||||
state_len >= seqlen,
|
||||
"causal_conv1d_update_cpu: state_len must be >= seqlen for 3D x.");
|
||||
TORCH_CHECK(
|
||||
conv_states.stride(-2) == 1 && conv_states.stride(-1) == dim,
|
||||
"causal_conv1d_update_cpu: 3D x requires SD conv_states layout.");
|
||||
|
||||
const int32_t* accepted_counts =
|
||||
num_accepted_tokens.value().data_ptr<int32_t>();
|
||||
const int32_t* indices = conv_state_indices.value().data_ptr<int32_t>();
|
||||
const int64_t num_slots = conv_states.size(0);
|
||||
for (int64_t bs = 0; bs < batch; ++bs) {
|
||||
const int32_t num_accepted = accepted_counts[bs];
|
||||
const int32_t conv_state_index = indices[bs];
|
||||
TORCH_CHECK(
|
||||
conv_state_index != pad_slot_id,
|
||||
"causal_conv1d_update_cpu: 3D x does not support pad slots.");
|
||||
TORCH_CHECK(
|
||||
conv_state_index >= 0 && conv_state_index < num_slots,
|
||||
"causal_conv1d_update_cpu: conv_state_indices out of range.");
|
||||
TORCH_CHECK(
|
||||
num_accepted >= 1 && num_accepted <= seqlen,
|
||||
"causal_conv1d_update_cpu: num_accepted_tokens must be in [1, "
|
||||
"seqlen].");
|
||||
TORCH_CHECK(
|
||||
num_accepted - 1 + width - 1 <= state_len,
|
||||
"causal_conv1d_update_cpu: history window exceeds conv_states.");
|
||||
}
|
||||
|
||||
int64_t conv_state_slot_stride = conv_states.stride(0);
|
||||
at::Tensor out = at::empty_like(x);
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(
|
||||
scalar_type, "causal_conv1d_update_multi_kernel_impl", [&] {
|
||||
causal_conv1d_update_multi_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
x.data_ptr<scalar_t>(),
|
||||
conv_states.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<scalar_t>(),
|
||||
conditional_data_ptr<scalar_t>(bias),
|
||||
accepted_counts,
|
||||
indices,
|
||||
silu_activation,
|
||||
batch,
|
||||
dim,
|
||||
seqlen,
|
||||
width,
|
||||
state_len,
|
||||
conv_state_slot_stride);
|
||||
});
|
||||
return out;
|
||||
}
|
||||
|
||||
TORCH_CHECK(
|
||||
!num_accepted_tokens.has_value(),
|
||||
"causal_conv1d_update_cpu: num_accepted_tokens is only supported for 3D "
|
||||
"x.");
|
||||
|
||||
// adjust `conv_states` to be contiguous on `dim`
|
||||
if (conv_states.stride(-2) != 1) {
|
||||
TORCH_CHECK(state_len == width - 1,
|
||||
"causal_conv1d_update_cpu: wide conv_states must be contiguous on dim.");
|
||||
int64_t num_cache_lines = conv_states.size(0);
|
||||
auto conv_states_copy = conv_states.clone();
|
||||
conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
|
||||
|
||||
@@ -147,7 +147,7 @@ at::Tensor causal_conv1d_fwd_cpu(
|
||||
at::Tensor causal_conv1d_update_cpu(
|
||||
const at::Tensor& x, const at::Tensor& conv_states,
|
||||
const at::Tensor& weight, const std::optional<at::Tensor>& bias,
|
||||
bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
|
||||
bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens,
|
||||
const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
|
||||
bool is_vnni);
|
||||
|
||||
@@ -207,6 +207,20 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
|
||||
const torch::Tensor& topk_id, const bool skip_weighted,
|
||||
const std::string& act, const std::string& isa);
|
||||
|
||||
void prepack_moe_weight_int8(const torch::Tensor& weight,
|
||||
torch::Tensor& packed_weight,
|
||||
const std::string& isa);
|
||||
|
||||
void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
|
||||
const torch::Tensor& w13, const torch::Tensor& w2,
|
||||
const torch::Tensor& w13_scale,
|
||||
const torch::Tensor& w2_scale,
|
||||
const std::optional<torch::Tensor>& w13_bias,
|
||||
const std::optional<torch::Tensor>& w2_bias,
|
||||
const torch::Tensor& topk_weights,
|
||||
const torch::Tensor& topk_id, const bool skip_weighted,
|
||||
const std::string& act, const std::string& isa);
|
||||
|
||||
void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
const torch::Tensor positions,
|
||||
const torch::Tensor block_table,
|
||||
@@ -502,7 +516,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
|
||||
"weight, Tensor? bias, bool silu_activation,"
|
||||
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
|
||||
"Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "
|
||||
"pad_slot_id, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
|
||||
#endif
|
||||
@@ -596,8 +611,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
#endif
|
||||
|
||||
// fused moe
|
||||
#if defined(__AVX512F__) || \
|
||||
(defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
|
||||
#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
|
||||
ops.def(
|
||||
"prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
|
||||
"-> ()");
|
||||
@@ -608,7 +622,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"bool skip_weighted, "
|
||||
"str act, str isa) -> ()");
|
||||
ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
|
||||
#endif
|
||||
#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) &&
|
||||
// !defined(__APPLE__))
|
||||
#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
|
||||
!defined(__APPLE__)
|
||||
ops.def(
|
||||
"prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, "
|
||||
"str isa) -> ()");
|
||||
ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8);
|
||||
ops.def(
|
||||
"cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, "
|
||||
"Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, "
|
||||
"Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool "
|
||||
"skip_weighted, str act, str isa) -> ()");
|
||||
ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8);
|
||||
#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) &&
|
||||
// !defined(__APPLE__)
|
||||
ops.def(
|
||||
"mla_decode_kvcache("
|
||||
" Tensor! out, Tensor query, Tensor kv_cache,"
|
||||
|
||||
@@ -22,17 +22,25 @@ template <typename AllReduceKernel, typename T>
|
||||
__global__ __quickreduce_launch_bounds_two_shot__ static void
|
||||
allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
|
||||
int rank, uint8_t** dbuffer_list,
|
||||
uint32_t data_offset, uint32_t flag_color,
|
||||
uint32_t data_offset, uint32_t* d_flag_counters,
|
||||
int64_t data_size_per_phase) {
|
||||
int block = blockIdx.x;
|
||||
int grid = gridDim.x;
|
||||
|
||||
// Load this block's counter from device memory and advance it on-device,
|
||||
// so the color keeps changing across graph replays instead of being frozen.
|
||||
uint32_t flag_color = d_flag_counters[blockIdx.x];
|
||||
|
||||
while (block < num_blocks) {
|
||||
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
|
||||
flag_color, data_size_per_phase);
|
||||
block += grid;
|
||||
flag_color++;
|
||||
}
|
||||
// All threads compute the same final value; one writer per block is enough.
|
||||
if (threadIdx.x == 0 && threadIdx.y == 0) {
|
||||
d_flag_counters[blockIdx.x] = flag_color;
|
||||
}
|
||||
}
|
||||
|
||||
#define TWOSHOT_DISPATCH(__codec) \
|
||||
@@ -42,21 +50,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
|
||||
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
|
||||
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
|
||||
num_blocks, rank, dbuffer_list, data_offset, \
|
||||
flag_color, this->kMaxProblemSize); \
|
||||
d_flag_counters, this->kMaxProblemSize); \
|
||||
} else if (world_size == 4) { \
|
||||
using LineCodec = __codec<T, 4>; \
|
||||
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
|
||||
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
|
||||
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
|
||||
num_blocks, rank, dbuffer_list, data_offset, \
|
||||
flag_color, this->kMaxProblemSize); \
|
||||
d_flag_counters, this->kMaxProblemSize); \
|
||||
} else if (world_size == 8) { \
|
||||
using LineCodec = __codec<T, 8>; \
|
||||
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
|
||||
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
|
||||
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
|
||||
num_blocks, rank, dbuffer_list, data_offset, \
|
||||
flag_color, this->kMaxProblemSize); \
|
||||
d_flag_counters, this->kMaxProblemSize); \
|
||||
}
|
||||
|
||||
// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
|
||||
@@ -69,7 +77,7 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
|
||||
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
|
||||
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
|
||||
num_blocks, rank, dbuffer_list, data_offset, \
|
||||
flag_color, this->kMaxProblemSize); \
|
||||
d_flag_counters, this->kMaxProblemSize); \
|
||||
} else { \
|
||||
throw std::runtime_error( \
|
||||
"INT3 quick all-reduce is only supported for world_size == 2 " \
|
||||
@@ -94,7 +102,7 @@ struct DeviceComms {
|
||||
static int constexpr kMaxWorldSize = 8;
|
||||
|
||||
bool initialized = false;
|
||||
uint32_t flag_color = 1;
|
||||
uint32_t* d_flag_counters = nullptr;
|
||||
int world_size;
|
||||
int rank;
|
||||
|
||||
@@ -128,6 +136,16 @@ struct DeviceComms {
|
||||
// Clear the flags buffer.
|
||||
HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
|
||||
|
||||
// One flag-color counter per block, advanced by the kernel. Start at 1
|
||||
// to stay clear of the flags buffer we just zeroed.
|
||||
HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
|
||||
{
|
||||
std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
|
||||
HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
|
||||
kMaxNumBlocks * sizeof(uint32_t),
|
||||
hipMemcpyHostToDevice));
|
||||
}
|
||||
|
||||
// Device-side list of IPC buffers.
|
||||
buffer_list.resize(world_size);
|
||||
HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
|
||||
@@ -144,6 +162,12 @@ struct DeviceComms {
|
||||
hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
|
||||
|
||||
void destroy() {
|
||||
// Allocated before `initialized` flips true, so free it on its own guard
|
||||
// to avoid a leak if init fails partway through.
|
||||
if (d_flag_counters) {
|
||||
HIP_CHECK(hipFree(d_flag_counters));
|
||||
d_flag_counters = nullptr;
|
||||
}
|
||||
if (initialized) {
|
||||
for (int i = 0; i < world_size; i++) {
|
||||
if (i != rank) {
|
||||
@@ -211,8 +235,6 @@ struct DeviceComms {
|
||||
break;
|
||||
}
|
||||
HIP_CHECK(cudaGetLastError());
|
||||
// Rotate the flag color.
|
||||
flag_color += divceil(N, grid);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
+25
-8
@@ -25,6 +25,10 @@
|
||||
ARG CUDA_VERSION=13.0.3
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
|
||||
# This version is only used for CUDA 13+ builds; CUDA 12 falls back to
|
||||
# the default NCCL version shipped with the base image.
|
||||
ARG NCCL_VERSION=2.30.7
|
||||
|
||||
# By parameterizing the base images, we allow third-party to use their own
|
||||
# base images. One use case is hermetic builds with base images stored in
|
||||
@@ -477,10 +481,17 @@ WORKDIR /workspace
|
||||
# Build DeepEP wheels
|
||||
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
|
||||
# Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
|
||||
ARG DEEPEP_COMMIT_HASH=73b6ea4
|
||||
ARG DEEPEP_COMMIT_HASH=d4f41e4e93
|
||||
ARG NVSHMEM_VER
|
||||
ARG NCCL_VERSION
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
mkdir -p /tmp/ep_kernels_workspace/dist && \
|
||||
CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
|
||||
if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \
|
||||
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
|
||||
> /tmp/nccl-override.txt && \
|
||||
export UV_OVERRIDE=/tmp/nccl-override.txt; \
|
||||
fi && \
|
||||
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
|
||||
/tmp/install_python_libraries.sh \
|
||||
--workspace /tmp/ep_kernels_workspace \
|
||||
@@ -644,6 +655,7 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
|
||||
|
||||
ARG CUDA_VERSION
|
||||
ARG PYTHON_VERSION
|
||||
ARG NCCL_VERSION
|
||||
ARG DEADSNAKES_MIRROR_URL
|
||||
ARG DEADSNAKES_GPGKEY_URL
|
||||
ARG GET_PIP_URL
|
||||
@@ -696,7 +708,6 @@ RUN apt-get update -y \
|
||||
# Install CUDA development tools for runtime JIT compilation
|
||||
# (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
|
||||
RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --allow-change-held-packages \
|
||||
cuda-nvcc-${CUDA_VERSION_DASH} \
|
||||
@@ -709,12 +720,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
libnuma-dev \
|
||||
# numactl CLI for NUMA binding at runtime
|
||||
numactl && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
|
||||
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
|
||||
# so we pin the version to match our CUDA version
|
||||
NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \
|
||||
apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install uv for faster pip installs
|
||||
@@ -734,6 +739,18 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
|
||||
&& chgrp -R 0 /opt/uv \
|
||||
&& chmod -R g+rwX,a+rX /opt/uv
|
||||
|
||||
# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins
|
||||
# an older version as a transitive dep; this override forces uv to use our
|
||||
# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op).
|
||||
RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
|
||||
if [ "$CUDA_MAJOR" -ge 13 ]; then \
|
||||
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
|
||||
> /etc/uv-overrides.txt; \
|
||||
else \
|
||||
touch /etc/uv-overrides.txt; \
|
||||
fi
|
||||
ENV UV_OVERRIDE=/etc/uv-overrides.txt
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Non-root support (opt-in)
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
@@ -745,6 +745,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
|
||||
|
||||
# ROCm profiler limits workaround.
|
||||
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
|
||||
@@ -835,6 +836,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
|
||||
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
|
||||
@@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
|
||||
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
|
||||
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
|
||||
ENV AITER_ROCM_ARCH=gfx942;gfx950
|
||||
ENV MORI_GPU_ARCHS=gfx942;gfx950
|
||||
# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime
|
||||
|
||||
# Required for RCCL in ROCm7.1
|
||||
ENV HSA_NO_SCRATCH_RECLAIM=1
|
||||
|
||||
@@ -10,6 +10,9 @@
|
||||
"UBUNTU_VERSION": {
|
||||
"default": "22.04"
|
||||
},
|
||||
"NCCL_VERSION": {
|
||||
"default": "2.30.7"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
|
||||
},
|
||||
@@ -56,7 +59,7 @@
|
||||
"default": "cuda"
|
||||
},
|
||||
"DEEPEP_COMMIT_HASH": {
|
||||
"default": "73b6ea4"
|
||||
"default": "d4f41e4e93"
|
||||
},
|
||||
"GIT_REPO_CHECK": {
|
||||
"default": "0"
|
||||
|
||||
+3
-1
@@ -56,7 +56,9 @@ nav:
|
||||
- API Reference:
|
||||
- api/README.md
|
||||
- api/vllm
|
||||
- CLI Reference: cli
|
||||
- CLI Reference:
|
||||
- cli/README.md
|
||||
- vllm: cli
|
||||
- Community:
|
||||
- community/*
|
||||
- Governance: governance
|
||||
|
||||
+7
-9
@@ -1,10 +1,8 @@
|
||||
nav:
|
||||
- README.md
|
||||
- serve.md
|
||||
- chat.md
|
||||
- complete.md
|
||||
- run-batch.md
|
||||
- vllm bench:
|
||||
- bench/**/*.md
|
||||
- vllm launch:
|
||||
- launch/**/*.md
|
||||
- "*.md"
|
||||
- bench:
|
||||
- bench/*.md
|
||||
- sweep:
|
||||
- bench/sweep/*.md
|
||||
- launch:
|
||||
- launch/*.md
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench latency
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_latency.inc.md"
|
||||
@@ -1,55 +0,0 @@
|
||||
# vllm bench mm-processor
|
||||
|
||||
## Overview
|
||||
|
||||
`vllm bench mm-processor` profiles the multimodal input processor pipeline of
|
||||
vision-language models. It measures per-stage latency from the HuggingFace
|
||||
processor through to the encoder forward pass, helping you identify
|
||||
preprocessing bottlenecks and understand how different image resolutions or
|
||||
item counts affect end-to-end request time.
|
||||
|
||||
The benchmark supports two data sources: synthetic random multimodal inputs
|
||||
(`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before
|
||||
measurement to ensure stable results.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
vllm bench mm-processor \
|
||||
--model Qwen/Qwen2-VL-7B-Instruct \
|
||||
--dataset-name random-mm \
|
||||
--num-prompts 50 \
|
||||
--random-input-len 300 \
|
||||
--random-output-len 40 \
|
||||
--random-mm-base-items-per-request 2 \
|
||||
--random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
|
||||
--random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}'
|
||||
```
|
||||
|
||||
## Measured Stages
|
||||
|
||||
| Stage | Description |
|
||||
| ----- | ----------- |
|
||||
| `get_mm_hashes_secs` | Time spent hashing multimodal inputs |
|
||||
| `get_cache_missing_items_secs` | Time spent looking up the processor cache |
|
||||
| `apply_hf_processor_secs` | Time spent in the HuggingFace processor |
|
||||
| `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs |
|
||||
| `apply_prompt_updates_secs` | Time spent updating prompt tokens |
|
||||
| `preprocessor_total_secs` | Total preprocessing time |
|
||||
| `encoder_forward_secs` | Time spent in the encoder model forward pass |
|
||||
| `num_encoder_calls` | Number of encoder invocations per request |
|
||||
|
||||
The benchmark also reports end-to-end latency (TTFT + decode time) per
|
||||
request. Use `--metric-percentiles` to select which percentiles to report
|
||||
(default: p99) and `--output-json` to save results.
|
||||
|
||||
For more examples (HF datasets, warmup, JSON output), see
|
||||
[Benchmarking CLI — Multimodal Processor Benchmark](../../benchmarking/cli.md#multimodal-processor-benchmark).
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_mm_processor.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench serve
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_serve.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench sweep plot
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_sweep_plot.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench sweep plot_pareto
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_sweep_plot_pareto.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench sweep serve
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_sweep_serve.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench sweep serve_workload
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_sweep_serve_workload.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm bench throughput
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/bench_throughput.inc.md"
|
||||
@@ -1,5 +0,0 @@
|
||||
# vllm chat
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/chat.inc.md"
|
||||
@@ -1,5 +0,0 @@
|
||||
# vllm complete
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/complete.inc.md"
|
||||
@@ -1,10 +0,0 @@
|
||||
<!-- markdownlint-disable MD041 -->
|
||||
When passing JSON CLI arguments, the following sets of arguments are equivalent:
|
||||
|
||||
- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'`
|
||||
- `--json-arg.key1 value1 --json-arg.key2.key3 value2`
|
||||
|
||||
Additionally, list elements can be passed individually using `+`:
|
||||
|
||||
- `--json-arg '{"key4": ["value3", "value4", "value5"]}'`
|
||||
- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'`
|
||||
@@ -1,22 +0,0 @@
|
||||
# vllm launch render
|
||||
|
||||
## Overview
|
||||
|
||||
`vllm launch render` starts a GPU-less rendering server for preprocessing and
|
||||
postprocessing only.
|
||||
|
||||
```bash
|
||||
vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100
|
||||
```
|
||||
|
||||
This command reuses the standard serving parser, so model, frontend,
|
||||
networking, and related CLI options follow the same conventions as
|
||||
[`vllm serve`](../serve.md).
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/launch_render.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm run-batch
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/run-batch.inc.md"
|
||||
@@ -1,9 +0,0 @@
|
||||
# vllm serve
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse/serve.inc.md"
|
||||
@@ -11,12 +11,4 @@ Engine arguments control the behavior of the vLLM engine.
|
||||
|
||||
The engine argument classes, [EngineArgs][vllm.engine.arg_utils.EngineArgs] and [AsyncEngineArgs][vllm.engine.arg_utils.AsyncEngineArgs], are a combination of the configuration classes defined in [vllm.config][]. Therefore, if you are interested in developer documentation, we recommend looking at these configuration classes as they are the source of truth for types, defaults and docstrings.
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
## `EngineArgs`
|
||||
|
||||
--8<-- "docs/generated/argparse/engine_args.inc.md"
|
||||
|
||||
## `AsyncEngineArgs`
|
||||
|
||||
--8<-- "docs/generated/argparse/async_engine_args.inc.md"
|
||||
--8<-- "gen:engine-args"
|
||||
|
||||
@@ -195,7 +195,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
|
||||
The API server takes care of basic audio I/O and optional chunking before building prompts:
|
||||
|
||||
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
|
||||
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
|
||||
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into chunks and generates a prompt per chunk. There is no overlap between chunks, overlap_chunk_second controls the size of the search window used to find the split point.
|
||||
- Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
|
||||
|
||||
Relevant server logic:
|
||||
|
||||
@@ -8,6 +8,26 @@ toc_depth: 2
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.md:pre-built-images"
|
||||
|
||||
## Persist the compile cache across containers
|
||||
|
||||
Mounting the Hugging Face cache keeps model weights across containers, but each
|
||||
new container still starts with an empty `VLLM_CACHE_ROOT` (default
|
||||
`~/.cache/vllm`) and recompiles the model's `torch.compile` artifacts. Mount a
|
||||
named volume at that path to reuse the inductor, Triton, and AOT artifacts from
|
||||
the second container onward:
|
||||
|
||||
```bash
|
||||
docker run --rm --gpus all \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
-v vllm-cache:/root/.cache/vllm \
|
||||
-p 8000:8000 \
|
||||
vllm/vllm-openai:latest \
|
||||
meta-llama/Llama-3.1-8B-Instruct
|
||||
```
|
||||
|
||||
See [Faster Startup](../configuration/optimization.md#faster-startup) for the
|
||||
mechanism and for what invalidates the cache.
|
||||
|
||||
## Run as a non-root user
|
||||
|
||||
The CUDA `vllm/vllm-openai` image runs as root by default for backward
|
||||
|
||||
@@ -1,15 +1,9 @@
|
||||
# Attention Backend Feature Support
|
||||
|
||||
This document is auto-generated by `tools/pre_commit/generate_attention_backend_docs.py`.
|
||||
It shows the feature support for each registered attention backend
|
||||
based on the checks in `AttentionBackend.validate_configuration()`.
|
||||
|
||||
**Do not edit this file manually.** Run the following command to
|
||||
regenerate it:
|
||||
|
||||
```bash
|
||||
python tools/pre_commit/generate_attention_backend_docs.py
|
||||
```
|
||||
The priority and feature tables on this page are auto-generated from the
|
||||
attention backend registry by
|
||||
`docs/mkdocs/gen_files/generate_attention_backends.py`, based on the checks in
|
||||
`AttentionBackend.validate_configuration()`.
|
||||
|
||||
## Setting the Attention Backend
|
||||
|
||||
@@ -98,40 +92,11 @@ Priority is **1 = highest** (tried first).
|
||||
|
||||
### Standard Attention (MHA, MQA, GQA)
|
||||
|
||||
**Blackwell (SM 10.x):**
|
||||
|
||||
| Priority | Backend |
|
||||
| -------- | ------- |
|
||||
| 1 | `FLASHINFER` |
|
||||
| 2 | `FLASH_ATTN` |
|
||||
| 3 | `TRITON_ATTN` |
|
||||
| 4 | `FLEX_ATTENTION` |
|
||||
| 5 | `TURBOQUANT` |
|
||||
|
||||
**Ampere/Hopper (SM 8.x-9.x):**
|
||||
|
||||
| Priority | Backend |
|
||||
| -------- | ------- |
|
||||
| 1 | `FLASH_ATTN` |
|
||||
| 2 | `FLASHINFER` |
|
||||
| 3 | `TRITON_ATTN` |
|
||||
| 4 | `FLEX_ATTENTION` |
|
||||
| 5 | `TURBOQUANT` |
|
||||
--8<-- "gen:priority-standard"
|
||||
|
||||
### MLA Attention (DeepSeek-style)
|
||||
|
||||
**Blackwell (SM 10.x):**
|
||||
|
||||
| Priority | Backend |
|
||||
| -------- | ------- |
|
||||
| 1 | `FLASHINFER_MLA` |
|
||||
| 2 | `TOKENSPEED_MLA` |
|
||||
| 3 | `CUTLASS_MLA` |
|
||||
| 4 | `FLASH_ATTN_MLA` |
|
||||
| 5 | `FLASHMLA` |
|
||||
| 6 | `TRITON_MLA` |
|
||||
| 7 | `FLASHINFER_MLA_SPARSE`**\*** |
|
||||
| 8 | `FLASHMLA_SPARSE` |
|
||||
--8<-- "gen:priority-mla"
|
||||
|
||||
> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
|
||||
>
|
||||
@@ -157,24 +122,7 @@ Priority is **1 = highest** (tried first).
|
||||
|
||||
## Standard Attention (MHA, MQA, GQA) Backends
|
||||
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
|
||||
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
|
||||
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
|
||||
| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
--8<-- "gen:table-standard"
|
||||
|
||||
> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
|
||||
>
|
||||
@@ -188,9 +136,7 @@ automatic priority lists above. A lightning indexer scores KV blocks, the
|
||||
top-k blocks (plus fixed init/local blocks) are selected, and attention
|
||||
attends only to those blocks; index keys live in a separate side cache.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
|
||||
| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
--8<-- "gen:table-minimax"
|
||||
|
||||
## MLA (Multi-head Latent Attention) Backends
|
||||
|
||||
@@ -203,38 +149,20 @@ To explicitly select a prefill backend, use
|
||||
Otherwise, the prefill backend is selected automatically at runtime based on
|
||||
hardware and configuration.
|
||||
|
||||
| Backend | Description | Dtypes | Compute Cap. | Notes |
|
||||
| ------- | ----------- | ------ | ------------ | ----- |
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
|
||||
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
|
||||
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
--8<-- "gen:table-mla-prefill"
|
||||
|
||||
> **‡** Automatic selection tries FlashAttention first. On Blackwell
|
||||
> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
|
||||
> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
|
||||
> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64,
|
||||
> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other
|
||||
> GPUs, only FlashAttention is considered.
|
||||
|
||||
### Decode Backends
|
||||
|
||||
MLA decode backends are selected using the standard
|
||||
`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `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 |
|
||||
--8<-- "gen:table-mla-decode"
|
||||
|
||||
### DeepSeek V4 Decode Backends
|
||||
|
||||
@@ -245,8 +173,4 @@ pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
|
||||
default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and
|
||||
`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x |
|
||||
| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
--8<-- "gen:table-mla-v4-decode"
|
||||
|
||||
@@ -129,6 +129,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
|
||||
| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
|
||||
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
|
||||
| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
|
||||
|
||||
@@ -122,8 +122,6 @@ For example:
|
||||
|
||||
--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"
|
||||
|
||||
--8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer"
|
||||
|
||||
--8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
|
||||
```
|
||||
|
||||
|
||||
@@ -242,4 +242,4 @@ See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2a
|
||||
|
||||
## FusedMoEExpertsModular
|
||||
|
||||
See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts.
|
||||
See [Fused MoE Kernel features](./moe_kernel_features.md#fused-experts-kernels) for a list of all the available modular experts.
|
||||
|
||||
@@ -47,7 +47,8 @@ sequenceDiagram
|
||||
else only one side present
|
||||
PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched
|
||||
end
|
||||
PWriter->>PWriter: ensure D handshake (one-time)
|
||||
PWriter->>PWriter: _ensure_handshake to D (async; defer WRITE)
|
||||
PWriter->>PWriter: handshake callback re-queues on _deferred_push_inbox, wake
|
||||
PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif
|
||||
|
||||
note over DWorker,DWriter: D side - completion accounting
|
||||
@@ -100,15 +101,20 @@ event:
|
||||
D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)
|
||||
even when there is no new metadata to act on.
|
||||
3. **Handshake-completion callback** (background handshake executor
|
||||
thread) — when a deferred D→P handshake finishes successfully, the
|
||||
future's done-callback re-enqueues the registration onto
|
||||
``_reg_send_inbox`` and sets the wake so the corresponding
|
||||
``send_notif`` runs on the writer (we never call ``send_notif`` from
|
||||
the executor thread). On this second pass ``_ensure_handshake``
|
||||
returns ``None`` (the agent is now connected), so the writer sends
|
||||
the ``PUSH_REG`` directly. If the handshake *failed*, the callback
|
||||
fails the request instead of re-enqueuing, so there is no retry
|
||||
loop.
|
||||
thread) — both handshakes run on the executor and never block the
|
||||
writer; their done-callbacks re-enqueue the deferred op and set the
|
||||
wake, since neither ``send_notif`` nor the NIXL WRITE may run off the
|
||||
writer thread:
|
||||
* the **D→P** handshake (before sending ``PUSH_REG``) re-enqueues the
|
||||
registration onto ``_reg_send_inbox``;
|
||||
* the **P→D** handshake (before a WRITE) re-enqueues the matched
|
||||
``(req_id, blocks, reg_data)`` onto ``_deferred_push_inbox``.
|
||||
|
||||
On this second pass ``_ensure_handshake`` returns ``None`` (the agent
|
||||
is now connected), so the writer sends the ``PUSH_REG`` / issues the
|
||||
WRITE directly. If a handshake *failed*, the callback fails or drops
|
||||
the request instead of re-enqueuing, so there is no retry loop (see
|
||||
Failure handling).
|
||||
|
||||
In addition to event-driven wakes, the writer self-polls at
|
||||
``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side
|
||||
@@ -227,6 +233,13 @@ Two per-request timers are armed on the scheduler:
|
||||
* **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —
|
||||
identical handling: ``_handle_failed_transfer`` marks the recv as
|
||||
failed.
|
||||
* **P-side handshake failure (P→D handshake before a WRITE)** — the
|
||||
future's done-callback logs ``push_handshake_failed`` and drops the
|
||||
request without re-queuing. It deliberately does *not* call
|
||||
``_handle_failed_transfer`` (there is no ``_recving_metadata`` entry to
|
||||
invalidate on the producer side, same reasoning as the WRITE-submission
|
||||
failure below). P's blocks are reclaimed by the ``_kv_lease_duration``
|
||||
lease and D's stale registration by its watchdog.
|
||||
* **P-side WRITE submission failure** — the WRITE handle (if any) is
|
||||
released and ``xfer_stats.record_failed_transfer()`` bumps the
|
||||
failure counter. We deliberately do not call
|
||||
@@ -246,10 +259,11 @@ existing NIXL connector:
|
||||
class — all subclasses of the existing base classes;
|
||||
* one dedicated background thread per worker;
|
||||
* a few cross-thread queues, each with a single consumer (the writer);
|
||||
most have one producer, except ``_reg_send_inbox``, which is fed both
|
||||
by the engine main thread (new registrations) and by the
|
||||
handshake-completion callback (registrations replayed after their
|
||||
D→P handshake finishes);
|
||||
most have one producer, except the two replay queues fed by both the
|
||||
engine main thread and a handshake-completion callback:
|
||||
``_reg_send_inbox`` (registrations replayed after their D→P handshake)
|
||||
and ``_deferred_push_inbox`` (matched pushes replayed after their P→D
|
||||
handshake);
|
||||
* one new notification type (`PUSH_REG:<msgpack>`).
|
||||
|
||||
Behavior on the engine main thread is otherwise unchanged. The writer
|
||||
|
||||
@@ -157,7 +157,7 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti
|
||||
|
||||
The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required.
|
||||
|
||||
PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes.
|
||||
The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected.
|
||||
|
||||
| Key | Required | Default | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
@@ -176,6 +176,71 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
|
||||
- `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide.
|
||||
- `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision.
|
||||
|
||||
#### Orchestration-Layer Protocol
|
||||
|
||||
The P2P tier does not decide *which* peer to pull from — that is the orchestration layer's job (the router/EPP and its scheduler). The orchestrator drives every transfer through a request's `kv_transfer_params` dict: it picks the request's role, allocates a unique transaction ID, and supplies the remote peer's address. All block lookup, hash matching, and NIXL transfer happen at the tier level below; the orchestrator only sets the correct role keys and enforces the allowed combinations.
|
||||
|
||||
Every vLLM instance is a symmetric **peer**. Per request it acts as a **consumer** (pulls KV blocks from a remote peer's CPU cache instead of computing locally) or a **producer** (serves blocks from its own CPU cache to remote consumers) — or both, on the same session, for different requests. Roles are chosen per request by the keys below; there are no fixed prefiller/decoder processes.
|
||||
|
||||
Three role keys are defined, each mapping to a sub-dict. All are optional; a request with none of them uses the tier only as a local CPU cache.
|
||||
|
||||
Each key names the **remote counterpart** this peer transfers with (not this
|
||||
peer's own role), so the name reads as "the remote ___ I transfer with".
|
||||
|
||||
| Key | Set on | Value fields | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `remote_decoder` | prefill producer request | `kv_request_id` | Peer computes KV and keeps it available in CPU cache for the remote decoder to pull. |
|
||||
| `remote_prefiller` | decode consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer pulls KV from the remote prefiller at the given address (classic P/D disaggregation). |
|
||||
| `remote_kv_source` | P2P consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer looks up and pulls whatever blocks the remote source currently holds in CPU cache. |
|
||||
|
||||
Field semantics:
|
||||
|
||||
- `kv_request_id` (str): unique transaction ID allocated by the orchestrator and pushed to every peer involved in the transfer; used to correlate the lookup, fetch, and transfer-done messages. The producer is implicit — it serves whatever block hashes it currently holds in its CPU cache for that ID.
|
||||
- `remote_host` (str): IP/hostname of the remote peer's control socket to query. Must be the peer's routable node IP (see [Environment Variables](#environment-variables)).
|
||||
- `remote_port` (int): the peer's bound control-socket port, i.e. `base + data_parallel_index` for the selected DP rank.
|
||||
|
||||
Allowed and forbidden combinations:
|
||||
|
||||
- **`remote_decoder` + `remote_kv_source`** is the only legal multi-key combination: a prefill producer may *also* act as a P2P consumer for the same request — skipping prefix prefill by pulling cached blocks from a source while still keeping its own computed blocks available for a downstream decoder.
|
||||
- Forbidden: `remote_prefiller` + `remote_decoder` (contradictory roles), `remote_prefiller` + `remote_kv_source` (two competing fetch sources), and all three together.
|
||||
|
||||
Minimal examples (values that would appear in the request's `kv_transfer_params`):
|
||||
|
||||
```python
|
||||
# Prefill producer — compute and keep KV for a remote decoder to pull
|
||||
kv_transfer_params = {"remote_decoder": {"kv_request_id": "<unique-transfer-id>"}}
|
||||
|
||||
# Decode consumer — pull KV from a specific prefiller (classic P/D)
|
||||
kv_transfer_params = {
|
||||
"remote_prefiller": {
|
||||
"kv_request_id": "<unique-transfer-id>",
|
||||
"remote_host": "<prefiller-node-ip>",
|
||||
"remote_port": 5710,
|
||||
}
|
||||
}
|
||||
|
||||
# P2P consumer — pull whatever the source already has cached
|
||||
kv_transfer_params = {
|
||||
"remote_kv_source": {
|
||||
"kv_request_id": "<unique-transfer-id>",
|
||||
"remote_host": "<source-node-ip>",
|
||||
"remote_port": 5710,
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the keys above:
|
||||
|
||||
1. Both peers already have listener threads on their control sockets (see [Environment Variables](#environment-variables)).
|
||||
2. **Lookup.** The consumer's tiering manager does per-block lookups; in P2P mode the tier returns `None` and registers the key. At `on_schedule_end` the consumer sends one **`LookupMsg`** (`kv_request_id` + block hashes) to the peer, per request, per step.
|
||||
3. The producer matches those hashes against its local CPU cache and replies with a **`LookupRespMsg`** carrying the hit block hashes.
|
||||
4. **Resolve.** Retried lookups now return hit / miss / in-flight. The consumer calls `submit_load` for hits only, allocating CPU slots only for hits.
|
||||
5. The consumer sends a **`FetchMsg`** (`kv_request_id`, block hashes, destination block indexes).
|
||||
6. The producer performs the **NIXL WRITE** transfer and sends **`TransferDone`** with a success status.
|
||||
7. On `get_finished`, hits are loaded into GPU as ordinary cache hits; misses are recomputed by the engine.
|
||||
|
||||
In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached.
|
||||
|
||||
## 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.
|
||||
|
||||
@@ -818,16 +818,18 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu
|
||||
|
||||
#### Video Decoding Backend
|
||||
|
||||
vLLM decodes video bytes into frames using a selectable decoding backend. Three
|
||||
vLLM decodes video bytes into frames using a selectable decoding backend. Five
|
||||
backends are supported:
|
||||
|
||||
- `opencv` (default): OpenCV-based decoder.
|
||||
- `pyav`: PyAV decoder.
|
||||
- `torchcodec`: TorchCodec (PyTorch-native) decoder.
|
||||
- `pynvvideocodec`: NVIDIA NVDEC-based decoder.
|
||||
- `deepstream`: NVIDIA DeepStream NVDEC-based decoder.
|
||||
|
||||
All three backends are ultimately backed by FFmpeg. `torchcodec` lets
|
||||
you choose which FFmpeg version is used while `opencv` and `pyav` rely on
|
||||
whichever FFmpeg build they were linked against.
|
||||
The CPU backends are backed by FFmpeg. `torchcodec` lets you choose which FFmpeg
|
||||
version is used while `opencv` and `pyav` rely on whichever FFmpeg build they
|
||||
were linked against.
|
||||
|
||||
Select the backend by passing the `backend` parameter via `--media-io-kwargs`:
|
||||
|
||||
@@ -854,6 +856,21 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
--media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}'
|
||||
```
|
||||
|
||||
**PyNvVideoCodec-specific parameters:**
|
||||
|
||||
- `hw_decoders`: Maximum number of concurrent hardware decoder slots retained
|
||||
by each API server process. It must be a positive integer and defaults to `2`,
|
||||
which is the recommended starting point for concurrent video workloads.
|
||||
Because vLLM reserves GPU memory for these slots at startup, this value cannot
|
||||
be overridden per request. Benchmark before increasing it because each
|
||||
additional slot increases the GPU memory reservation.
|
||||
|
||||
```bash
|
||||
# Example: explicitly use the recommended 2 hardware decoders
|
||||
vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
--media-io-kwargs '{"video": {"backend": "pynvvideocodec", "hw_decoders": 2}}'
|
||||
```
|
||||
|
||||
#### Video Frame Recovery
|
||||
|
||||
For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place.
|
||||
|
||||
@@ -19,6 +19,20 @@ following `quantization.quant_algo` values:
|
||||
- `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`).
|
||||
- `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`).
|
||||
|
||||
!!! note
|
||||
For NVFP4 checkpoints, vLLM selects a GEMM kernel automatically at load
|
||||
time from the backends available on the current platform (CUTLASS,
|
||||
FlashInfer, Marlin, and others). On GPUs without a supported native FP4
|
||||
GEMM kernel, vLLM falls back to weight-only (W4A16) execution via Marlin
|
||||
and logs a warning; this may reduce throughput for compute-heavy
|
||||
workloads. Use `--linear-backend` to override the automatic selection
|
||||
(this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment
|
||||
variable). Values relevant to NVFP4 include `cutlass`,
|
||||
`flashinfer_cutlass`, `flashinfer_trtllm`, `flashinfer_cudnn`, and
|
||||
`marlin`; the full list is documented under `KernelConfig` on the
|
||||
[Engine Arguments](../../configuration/engine_args.md) page and shown by
|
||||
`vllm serve --help=KernelConfig`.
|
||||
|
||||
## Quantizing HuggingFace Models with PTQ
|
||||
|
||||
You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory.
|
||||
|
||||
+183
-48
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import importlib.metadata
|
||||
import importlib.util
|
||||
import inspect
|
||||
import logging
|
||||
import sys
|
||||
import textwrap
|
||||
@@ -10,17 +11,21 @@ from argparse import SUPPRESS, Action, HelpFormatter
|
||||
from collections.abc import Callable, Iterable
|
||||
from importlib.machinery import ModuleSpec
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import mkdocs_gen_files
|
||||
import regex as re
|
||||
from pydantic_core import core_schema
|
||||
|
||||
logger = logging.getLogger("mkdocs")
|
||||
|
||||
ROOT_DIR = Path(__file__).parent.parent.parent.parent
|
||||
ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse"
|
||||
|
||||
sys.path.insert(0, str(ROOT_DIR))
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from generated_content import fill_markers # noqa: E402
|
||||
|
||||
|
||||
def mock_if_no_torch(mock_module: str, mock: MagicMock):
|
||||
@@ -132,8 +137,8 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100):
|
||||
|
||||
|
||||
bench_latency = auto_mock("vllm.benchmarks", "latency")
|
||||
bench_mm_processor = auto_mock("vllm.benchmarks", "mm_processor")
|
||||
bench_serve = auto_mock("vllm.benchmarks", "serve")
|
||||
bench_startup = auto_mock("vllm.benchmarks", "startup")
|
||||
bench_sweep_plot = auto_mock("vllm.benchmarks.sweep.plot", "SweepPlotArgs")
|
||||
bench_sweep_plot_pareto = auto_mock(
|
||||
"vllm.benchmarks.sweep.plot_pareto", "SweepPlotParetoArgs"
|
||||
@@ -142,12 +147,28 @@ bench_sweep_serve = auto_mock("vllm.benchmarks.sweep.serve", "SweepServeArgs")
|
||||
bench_sweep_serve_workload = auto_mock(
|
||||
"vllm.benchmarks.sweep.serve_workload", "SweepServeWorkloadArgs"
|
||||
)
|
||||
bench_sweep_startup = auto_mock("vllm.benchmarks.sweep.startup", "SweepStartupArgs")
|
||||
bench_throughput = auto_mock("vllm.benchmarks", "throughput")
|
||||
AsyncEngineArgs = auto_mock("vllm.engine.arg_utils", "AsyncEngineArgs")
|
||||
EngineArgs = auto_mock("vllm.engine.arg_utils", "EngineArgs")
|
||||
ChatCommand = auto_mock("vllm.entrypoints.cli.openai", "ChatCommand")
|
||||
CompleteCommand = auto_mock("vllm.entrypoints.cli.openai", "CompleteCommand")
|
||||
BenchmarkSubcommand = auto_mock(
|
||||
"vllm.entrypoints.cli.benchmark.main", "BenchmarkSubcommand"
|
||||
)
|
||||
import_bench_subcommands = auto_mock(
|
||||
"vllm.entrypoints.cli.benchmark.main", "_import_bench_subcommand_modules"
|
||||
)
|
||||
BenchmarkSubcommandBase = auto_mock(
|
||||
"vllm.entrypoints.cli.benchmark.base", "BenchmarkSubcommandBase"
|
||||
)
|
||||
BenchmarkMMProcessorSubcommand = auto_mock(
|
||||
"vllm.entrypoints.cli.benchmark.mm_processor", "BenchmarkMMProcessorSubcommand"
|
||||
)
|
||||
LaunchSubcommandBase = auto_mock("vllm.entrypoints.cli.launch", "LaunchSubcommandBase")
|
||||
launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION")
|
||||
RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand")
|
||||
sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS")
|
||||
openai_cli_args = auto_mock("vllm.entrypoints.openai", "cli_args")
|
||||
openai_run_batch = auto_mock("vllm.entrypoints.openai", "run_batch")
|
||||
|
||||
@@ -179,7 +200,7 @@ class MarkdownFormatter(HelpFormatter):
|
||||
|
||||
def add_text(self, text: str):
|
||||
if text:
|
||||
self._markdown_output.append(f"{text.strip()}\n\n")
|
||||
self._markdown_output.append(f"{inspect.cleandoc(text)}\n\n")
|
||||
|
||||
def add_usage(self, usage, actions, groups, prefix=None):
|
||||
pass
|
||||
@@ -241,49 +262,163 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser:
|
||||
return _parser or parser
|
||||
|
||||
|
||||
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
|
||||
logger.info("Generating argparse documentation")
|
||||
logger.debug("Root directory: %s", ROOT_DIR.resolve())
|
||||
logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
|
||||
|
||||
# Create the ARGPARSE_DOC_DIR if it doesn't exist
|
||||
if not ARGPARSE_DOC_DIR.exists():
|
||||
ARGPARSE_DOC_DIR.mkdir(parents=True)
|
||||
|
||||
# Create parsers to document
|
||||
parsers = {
|
||||
# Engine args
|
||||
"engine_args": create_parser(EngineArgs.add_cli_args),
|
||||
"async_engine_args": create_parser(
|
||||
AsyncEngineArgs.add_cli_args, async_args_only=True
|
||||
),
|
||||
# CLI
|
||||
"serve": create_parser(openai_cli_args.make_arg_parser),
|
||||
"chat": create_parser(ChatCommand.add_cli_args),
|
||||
"complete": create_parser(CompleteCommand.add_cli_args),
|
||||
"launch_render": create_parser(RenderSubcommand.add_cli_args),
|
||||
"run-batch": create_parser(openai_run_batch.make_arg_parser),
|
||||
# Benchmark CLI
|
||||
"bench_latency": create_parser(bench_latency.add_cli_args),
|
||||
"bench_mm_processor": create_parser(bench_mm_processor.add_cli_args),
|
||||
"bench_serve": create_parser(bench_serve.add_cli_args),
|
||||
"bench_sweep_plot": create_parser(bench_sweep_plot.add_cli_args),
|
||||
"bench_sweep_plot_pareto": create_parser(bench_sweep_plot_pareto.add_cli_args),
|
||||
"bench_sweep_serve": create_parser(bench_sweep_serve.add_cli_args),
|
||||
"bench_sweep_serve_workload": create_parser(
|
||||
bench_sweep_serve_workload.add_cli_args
|
||||
),
|
||||
"bench_throughput": create_parser(bench_throughput.add_cli_args),
|
||||
}
|
||||
|
||||
# Generate documentation for each parser
|
||||
for stem, parser in parsers.items():
|
||||
doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
|
||||
# Specify encoding for building on Windows
|
||||
with open(doc_path, "w", encoding="utf-8") as f:
|
||||
f.write(super(type(parser), parser).format_help())
|
||||
logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
|
||||
def format_help(parser: FlexibleArgumentParser) -> str:
|
||||
"""Format a parser's help as markdown using `MarkdownFormatter`."""
|
||||
return super(type(parser), parser).format_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
on_startup("build", False)
|
||||
# Absolute docs URLs are kept in the help text because they are useful in the
|
||||
# terminal. Wrap them as markdown links so the `url_schemes` hook can rewrite
|
||||
# them into doc-relative links / cross-references at render time.
|
||||
_DOCS_URL = re.compile(r"https://docs\.vllm\.ai/en/[^/\s]+/[^\s)>]+")
|
||||
|
||||
|
||||
def linkify_docs_urls(text: str) -> str:
|
||||
"""Wrap bare docs.vllm.ai URLs in help text as markdown links."""
|
||||
return _DOCS_URL.sub(lambda m: f"[{m.group()}]({m.group()})", text)
|
||||
|
||||
|
||||
logger.info("Generating argparse documentation")
|
||||
logger.debug("Root directory: %s", ROOT_DIR.resolve())
|
||||
|
||||
# The JSON tip is always rendered immediately before generated argument content,
|
||||
# and the generator is its only consumer, so it lives here rather than in a
|
||||
# separate snippet file. (The runtime terminal equivalent is
|
||||
# `FlexibleArgumentParser._json_tip` in vllm/utils/argparse_utils.py.)
|
||||
JSON_TIP = """## JSON CLI Arguments
|
||||
|
||||
When passing JSON CLI arguments, the following sets of arguments are equivalent:
|
||||
|
||||
- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'`
|
||||
- `--json-arg.key1 value1 --json-arg.key2.key3 value2`
|
||||
|
||||
Additionally, list elements can be passed individually using `+`:
|
||||
|
||||
- `--json-arg '{"key4": ["value3", "value4", "value5"]}'`
|
||||
- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'`
|
||||
|
||||
"""
|
||||
|
||||
# Argument sections filled into `gen:` markers on handwritten pages
|
||||
engine_args = create_parser(EngineArgs.add_cli_args)
|
||||
async_engine_args = create_parser(AsyncEngineArgs.add_cli_args, async_args_only=True)
|
||||
fill_markers(
|
||||
"configuration/engine_args.md",
|
||||
{
|
||||
"engine-args": (
|
||||
f"{JSON_TIP}## `EngineArgs`\n\n"
|
||||
f"{linkify_docs_urls(format_help(engine_args))}"
|
||||
f"## `AsyncEngineArgs`\n\n"
|
||||
f"{linkify_docs_urls(format_help(async_engine_args))}"
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
# CLI reference pages generated entirely from their parser: page -> (parser, JSON tip)
|
||||
pages = {
|
||||
"cli/serve.md": (create_parser(openai_cli_args.make_arg_parser), True),
|
||||
"cli/chat.md": (create_parser(ChatCommand.add_cli_args), False),
|
||||
"cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False),
|
||||
"cli/run-batch.md": (create_parser(openai_run_batch.make_arg_parser), True),
|
||||
"cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True),
|
||||
"cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True),
|
||||
# URL kept as `mm_processor` for back-compat; command name is `mm-processor`
|
||||
"cli/bench/mm_processor.md": (
|
||||
create_parser(BenchmarkMMProcessorSubcommand.add_cli_args),
|
||||
True,
|
||||
),
|
||||
"cli/bench/serve.md": (create_parser(bench_serve.add_cli_args), True),
|
||||
"cli/bench/startup.md": (create_parser(bench_startup.add_cli_args), True),
|
||||
"cli/bench/throughput.md": (create_parser(bench_throughput.add_cli_args), True),
|
||||
"cli/bench/sweep/plot.md": (create_parser(bench_sweep_plot.add_cli_args), True),
|
||||
"cli/bench/sweep/plot_pareto.md": (
|
||||
create_parser(bench_sweep_plot_pareto.add_cli_args),
|
||||
True,
|
||||
),
|
||||
"cli/bench/sweep/serve.md": (create_parser(bench_sweep_serve.add_cli_args), True),
|
||||
"cli/bench/sweep/serve_workload.md": (
|
||||
create_parser(bench_sweep_serve_workload.add_cli_args),
|
||||
True,
|
||||
),
|
||||
"cli/bench/sweep/startup.md": (
|
||||
create_parser(bench_sweep_startup.add_cli_args),
|
||||
True,
|
||||
),
|
||||
}
|
||||
|
||||
# Command name for pages whose file stem differs (URL kept for back-compat).
|
||||
COMMAND_NAMES = {"cli/bench/mm_processor.md": "mm-processor"}
|
||||
|
||||
for doc_path, (parser, json_tip) in pages.items():
|
||||
segments = Path(doc_path).relative_to("cli").with_suffix("").parts
|
||||
label = COMMAND_NAMES.get(doc_path, segments[-1])
|
||||
command = " ".join([*segments[:-1], label])
|
||||
# `title` frontmatter keeps the nav label to just this command's segment,
|
||||
# while the H1 stays the full `vllm ...` command for the page heading.
|
||||
content = f"---\ntitle: {label}\n---\n\n"
|
||||
content += f"# vllm {command}\n\n"
|
||||
if parser.description:
|
||||
content += f"## Overview\n\n{parser.description}\n\n"
|
||||
# Rendered above instead of at the top of the Arguments section
|
||||
parser.description = None
|
||||
if json_tip:
|
||||
content += JSON_TIP
|
||||
content += f"## Arguments\n\n{linkify_docs_urls(format_help(parser))}"
|
||||
with mkdocs_gen_files.open(doc_path, "w") as f:
|
||||
f.write(content)
|
||||
logger.debug("CLI reference generated: %s", doc_path)
|
||||
|
||||
logger.info("Total argparse docs generated: %d", len(pages) + 2)
|
||||
|
||||
|
||||
# --- Bare subcommand (group) pages -------------------------------------------
|
||||
# Mirror `vllm <group> --help`: an overview plus a table of child subcommands,
|
||||
# each linked to its reference page. Children are read from the CLI registries
|
||||
# so the listing can never drift from the actual subcommands. Each page is the
|
||||
# `README.md` of its command directory so it becomes that section's index and is
|
||||
# picked up by the existing nav globs.
|
||||
import_bench_subcommands() # populate BenchmarkSubcommandBase.__subclasses__()
|
||||
bench_subcommands = BenchmarkSubcommandBase.__subclasses__()
|
||||
bench_children = [(cmd.name, cmd.help) for cmd in bench_subcommands]
|
||||
|
||||
groups = {
|
||||
"cli/bench/README.md": (BenchmarkSubcommand.help, bench_children),
|
||||
"cli/launch/README.md": (
|
||||
launch_description,
|
||||
[(cmd.name, cmd.help) for cmd in LaunchSubcommandBase.__subclasses__()],
|
||||
),
|
||||
"cli/bench/sweep/README.md": (
|
||||
dict(bench_children).get("sweep"),
|
||||
[(args.parser_name, args.parser_help) for args, _ in sweep_subcommands],
|
||||
),
|
||||
}
|
||||
|
||||
# Doc paths that exist, so we only link a child that has a reference page.
|
||||
existing_pages = set(pages) | set(groups)
|
||||
|
||||
|
||||
def child_link(group_doc: str, name: str) -> str | None:
|
||||
group_dir = Path(group_doc).parent # cli/bench/README.md -> cli/bench
|
||||
for stem in (name, name.replace("-", "_")):
|
||||
# A leaf page (bench/latency.md) or a nested group index (sweep/README.md)
|
||||
for candidate in (group_dir / f"{stem}.md", group_dir / stem / "README.md"):
|
||||
if candidate.as_posix() in existing_pages:
|
||||
return candidate.relative_to(group_dir).as_posix()
|
||||
return None
|
||||
|
||||
|
||||
for doc_path, (overview, children) in groups.items():
|
||||
title = "vllm " + Path(doc_path).parent.relative_to("cli").as_posix()
|
||||
lines = [f"# {title.replace('/', ' ')}", ""]
|
||||
if overview:
|
||||
lines += ["## Overview", "", overview.strip(), ""]
|
||||
lines += ["## Subcommands", "", "| Command | Description |", "| --- | --- |"]
|
||||
for name, summary in children:
|
||||
link = child_link(doc_path, name)
|
||||
command = f"[`{name}`]({link})" if link else f"`{name}`"
|
||||
lines.append(f"| {command} | {(summary or '').strip()} |")
|
||||
with mkdocs_gen_files.open(doc_path, "w") as f:
|
||||
f.write("\n".join(lines) + "\n")
|
||||
logger.debug("CLI group reference generated: %s", doc_path)
|
||||
|
||||
logger.info("CLI group reference pages generated: %d", len(groups))
|
||||
+61
-404
@@ -9,33 +9,28 @@ based on the checks in AttentionBackend.validate_configuration().
|
||||
|
||||
This approach avoids requiring CUDA/ROCm/GPU libraries to be installed.
|
||||
|
||||
When used as a pre-commit hook, this script receives filenames as arguments
|
||||
and only runs the check if any of the relevant files were modified.
|
||||
It runs as an mkdocs-gen-files script, so the page is generated at docs build
|
||||
time rather than being committed to the repository.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import fnmatch
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from generated_content import fill_markers # noqa: E402
|
||||
|
||||
logger = logging.getLogger("mkdocs")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants and file paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
REPO_ROOT = Path(__file__).parent.parent.parent
|
||||
|
||||
RELEVANT_PATTERNS = [
|
||||
"vllm/v1/attention/backends/*.py",
|
||||
"vllm/v1/attention/backends/**/*.py",
|
||||
"vllm/models/minimax_m3/common/sparse_attention.py",
|
||||
"vllm/model_executor/layers/attention/mla_attention.py",
|
||||
"vllm/platforms/cuda.py",
|
||||
"tools/pre_commit/generate_attention_backend_docs.py",
|
||||
"docs/design/attention_backends.md",
|
||||
]
|
||||
REPO_ROOT = Path(__file__).parent.parent.parent.parent
|
||||
|
||||
BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends"
|
||||
REGISTRY_FILE = BACKENDS_DIR / "registry.py"
|
||||
@@ -55,19 +50,6 @@ BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = {
|
||||
}
|
||||
|
||||
|
||||
def is_relevant_file(filepath: str) -> bool:
|
||||
"""Check if a file matches any of the relevant patterns."""
|
||||
path = Path(filepath)
|
||||
if path.is_absolute():
|
||||
try:
|
||||
path = path.relative_to(REPO_ROOT)
|
||||
except ValueError:
|
||||
return False
|
||||
path_str = str(path)
|
||||
|
||||
return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS)
|
||||
|
||||
|
||||
MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill"
|
||||
MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py"
|
||||
MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py"
|
||||
@@ -960,7 +942,7 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None
|
||||
try:
|
||||
tree = ast.parse(file_path.read_text())
|
||||
except Exception as e:
|
||||
print(f" Warning: Could not parse {file_path}: {e}", file=sys.stderr)
|
||||
logger.warning("Could not parse %s: %s", file_path, e)
|
||||
return None
|
||||
|
||||
class_name = class_path.rsplit(".", 1)[1]
|
||||
@@ -1657,113 +1639,12 @@ def _render_table(
|
||||
return lines
|
||||
|
||||
|
||||
def generate_markdown_table(
|
||||
backends: list[dict[str, Any]], title: str, is_mla_table: bool = False
|
||||
) -> str:
|
||||
"""Generate a titled markdown table from backend info."""
|
||||
if not backends:
|
||||
return f"## {title}\n\nNo backends found.\n"
|
||||
has_versions = any(b.get("version") for b in backends)
|
||||
columns = _build_columns(is_mla_table, has_versions)
|
||||
lines = [f"## {title}", ""]
|
||||
lines.extend(_render_table(columns, backends))
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Markdown section generators (usage, priority, legend, MLA)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def generate_usage_section() -> str:
|
||||
"""Generate the usage documentation section."""
|
||||
return """## Setting the Attention Backend
|
||||
|
||||
### Command Line
|
||||
|
||||
There are two ways to specify the backend from the command line:
|
||||
|
||||
**Option 1: Using `--attention-backend` (simple)**
|
||||
|
||||
```bash
|
||||
vllm serve <model> --attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
**Option 2: Using `--attention-config.backend` / `-ac.backend` (structured config)**
|
||||
|
||||
```bash
|
||||
# Dot notation
|
||||
vllm serve <model> --attention-config.backend FLASH_ATTN
|
||||
vllm serve <model> -ac.backend FLASH_ATTN
|
||||
|
||||
# JSON format
|
||||
vllm serve <model> --attention-config '{"backend": "FLASH_ATTN"}'
|
||||
vllm serve <model> -ac '{"backend": "FLASH_ATTN"}'
|
||||
```
|
||||
|
||||
> **Note:** `--attention-backend` and `--attention-config.backend` are mutually
|
||||
> exclusive. Use one or the other, not both.
|
||||
|
||||
### Python API
|
||||
|
||||
Use `AttentionConfig` with the `LLM` class:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.config import AttentionConfig
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
# Method 1: Using AttentionConfig with enum
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
attention_config=AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN),
|
||||
)
|
||||
|
||||
# Method 2: Using attention_backend parameter with string
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
attention_backend="FLASH_ATTN",
|
||||
)
|
||||
```
|
||||
|
||||
## Backend Selection Behavior
|
||||
|
||||
### Manual Selection
|
||||
|
||||
When you explicitly set a backend via `--attention-backend` or `AttentionConfig`:
|
||||
|
||||
1. The backend is **validated** against your configuration (model dtype, head
|
||||
size, compute capability, etc.)
|
||||
2. If the backend **doesn't support** your configuration, an error is raised
|
||||
with the specific reason
|
||||
3. If valid, the backend is used
|
||||
|
||||
Example error when selecting an incompatible backend:
|
||||
|
||||
```text
|
||||
ValueError: Selected backend FLASHMLA is not valid for this configuration.
|
||||
Reason: ['compute capability not supported']
|
||||
```
|
||||
|
||||
### Automatic Selection
|
||||
|
||||
When no backend is specified (the default):
|
||||
|
||||
1. vLLM iterates through backends in **priority order** (see tables below)
|
||||
2. Each backend is validated against your configuration
|
||||
3. The **first compatible backend** is selected
|
||||
4. If no backend is compatible, an error is raised listing all backends and
|
||||
their incompatibility reasons
|
||||
"""
|
||||
|
||||
|
||||
def _priority_table(
|
||||
title: str,
|
||||
backends: list[str],
|
||||
annotations: dict[str, str] | None = None,
|
||||
) -> list[str]:
|
||||
"""Generate a priority table for a list of backends."""
|
||||
"""Render a priority table for a list of backends."""
|
||||
|
||||
def _fmt(b: str) -> str:
|
||||
suffix = annotations.get(b, "") if annotations else ""
|
||||
@@ -1779,102 +1660,38 @@ def _priority_table(
|
||||
]
|
||||
|
||||
|
||||
def generate_priority_section(priorities: dict[str, list[str]]) -> str:
|
||||
"""Generate the priority ranking section."""
|
||||
lines = [
|
||||
"## Backend Priority (CUDA)",
|
||||
"",
|
||||
"When no backend is explicitly selected, vLLM chooses the first",
|
||||
"compatible backend from these priority-ordered lists.",
|
||||
"",
|
||||
"Priority is **1 = highest** (tried first).",
|
||||
"",
|
||||
"### Standard Attention (MHA, MQA, GQA)",
|
||||
"",
|
||||
]
|
||||
|
||||
sm100 = "Blackwell (SM 10.x)"
|
||||
ampere = "Ampere/Hopper (SM 8.x-9.x)"
|
||||
|
||||
if "standard_sm100" in priorities:
|
||||
lines.extend(_priority_table(sm100, priorities["standard_sm100"]))
|
||||
if "standard_default" in priorities:
|
||||
lines.extend(_priority_table(ampere, priorities["standard_default"]))
|
||||
|
||||
lines.extend(["### MLA Attention (DeepSeek-style)", ""])
|
||||
|
||||
mla_sm100_annotations = {
|
||||
"FLASHINFER_MLA_SPARSE": "**\\***",
|
||||
}
|
||||
if "mla_sm100" in priorities:
|
||||
lines.extend(
|
||||
_priority_table(sm100, priorities["mla_sm100"], mla_sm100_annotations)
|
||||
)
|
||||
if "mla_default" in priorities:
|
||||
lines.extend(_priority_table(ampere, priorities["mla_default"]))
|
||||
|
||||
if "mla_sm100" in priorities:
|
||||
lines.append(
|
||||
"> **\\*** For sparse MLA, FP8 KV cache always prefers "
|
||||
"`FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` "
|
||||
"is preferred for low query-head counts (<= 16), while "
|
||||
"`FLASHMLA_SPARSE` is preferred otherwise."
|
||||
)
|
||||
lines.append(">")
|
||||
|
||||
lines.append(
|
||||
"> **Note:** ROCm and CPU platforms have their own selection logic. "
|
||||
"See the platform-specific documentation for details."
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
return "\n".join(lines)
|
||||
_SM100 = "Blackwell (SM 10.x)"
|
||||
_AMPERE = "Ampere/Hopper (SM 8.x-9.x)"
|
||||
|
||||
|
||||
def generate_legend() -> str:
|
||||
"""Generate a legend explaining the table columns."""
|
||||
return """## Legend
|
||||
|
||||
| Column | Description |
|
||||
| ------ | ----------- |
|
||||
| **Dtypes** | Supported model data types (fp16, bf16, fp32) |
|
||||
| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) |
|
||||
| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
|
||||
| **Head Sizes** | Supported attention head sizes |
|
||||
| **Sink** | Attention sink support (for StreamingLLM) |
|
||||
| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models |
|
||||
| **Sparse** | Sparse attention support (MLA only) |
|
||||
| **MM Prefix** | Multimodal prefix full attention support |
|
||||
| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
|
||||
| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
|
||||
| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
|
||||
|
||||
**Symbols:** ✅ = Supported, ❌ = Not supported
|
||||
"""
|
||||
|
||||
|
||||
def generate_mla_section(
|
||||
prefill_backends: list[dict[str, Any]],
|
||||
decode_backends: list[dict[str, Any]],
|
||||
v4_decode_backends: list[dict[str, Any]] | None = None,
|
||||
def _priority_block(
|
||||
priorities: dict[str, list[str]],
|
||||
sm100_key: str,
|
||||
default_key: str,
|
||||
sm100_annotations: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""Generate the complete MLA section with prefill and decode tables."""
|
||||
"""Render whichever priority tables exist for one attention category."""
|
||||
lines: list[str] = []
|
||||
if sm100_key in priorities:
|
||||
lines += _priority_table(_SM100, priorities[sm100_key], sm100_annotations)
|
||||
if default_key in priorities:
|
||||
lines += _priority_table(_AMPERE, priorities[default_key])
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def _feature_table(backends: list[dict[str, Any]], is_mla: bool) -> str:
|
||||
"""Render a backend feature table (header, separator, one row per backend)."""
|
||||
has_versions = any(b.get("version") for b in backends)
|
||||
columns = _build_columns(is_mla, has_versions)
|
||||
return "\n".join(_render_table(columns, backends))
|
||||
|
||||
|
||||
def _mla_prefill_table(prefill_backends: list[dict[str, Any]]) -> str:
|
||||
"""Render the MLA prefill backend table."""
|
||||
lines = [
|
||||
"## MLA (Multi-head Latent Attention) Backends",
|
||||
"",
|
||||
"MLA uses separate backends for prefill and decode phases.",
|
||||
"",
|
||||
"### Prefill Backends",
|
||||
"",
|
||||
"To explicitly select a prefill backend, use",
|
||||
"`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).",
|
||||
"Otherwise, the prefill backend is selected automatically at runtime based on",
|
||||
"hardware and configuration.",
|
||||
"",
|
||||
"| Backend | Description | Dtypes | Compute Cap. | Notes |",
|
||||
"| ------- | ----------- | ------ | ------------ | ----- |",
|
||||
]
|
||||
|
||||
for backend in prefill_backends:
|
||||
row = "| `{}`{} | {} | {} | {} | {} |".format(
|
||||
backend["name"],
|
||||
@@ -1885,87 +1702,21 @@ def generate_mla_section(
|
||||
backend.get("notes", ""),
|
||||
)
|
||||
lines.append(row.replace(" ", " "))
|
||||
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
"> **‡** Automatic selection tries FlashAttention first. On Blackwell",
|
||||
"> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then",
|
||||
"> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.",
|
||||
"",
|
||||
"### Decode Backends",
|
||||
"",
|
||||
"MLA decode backends are selected using the standard",
|
||||
"`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
# Reuse data-driven table rendering for decode backends
|
||||
columns = _build_columns(is_mla=True, has_versions=False)
|
||||
lines.extend(_render_table(columns, decode_backends))
|
||||
|
||||
if v4_decode_backends:
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
"### 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 `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and",
|
||||
"`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.",
|
||||
"",
|
||||
]
|
||||
)
|
||||
lines.extend(_render_table(columns, v4_decode_backends))
|
||||
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def generate_minimax_section(backends: list[dict[str, Any]]) -> str:
|
||||
"""Generate the MiniMax M3 sparse attention section."""
|
||||
lines = [
|
||||
"## MiniMax M3 Sparse Attention Backends",
|
||||
"",
|
||||
'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")',
|
||||
"layers. It is wired in directly by the model and is not part of the",
|
||||
"automatic priority lists above. A lightning indexer scores KV blocks, the",
|
||||
"top-k blocks (plus fixed init/local blocks) are selected, and attention",
|
||||
"attends only to those blocks; index keys live in a separate side cache.",
|
||||
"",
|
||||
]
|
||||
columns = _build_columns(is_mla=False, has_versions=False)
|
||||
lines.extend(_render_table(columns, backends))
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
def build_blocks() -> dict[str, str]:
|
||||
"""Build the generated table blocks keyed by their `gen:` marker name.
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Top-level orchestration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def generate_docs() -> str:
|
||||
"""Generate the complete documentation."""
|
||||
Only the tables are generated here; the surrounding prose lives in the
|
||||
handwritten ``docs/design/attention_backends.md`` page.
|
||||
"""
|
||||
attention_backends_map = parse_registry()
|
||||
|
||||
# Parse priority lists from cuda.py
|
||||
priorities = parse_cuda_priority_lists()
|
||||
|
||||
# Parse FlashAttention FA2/FA3 feature differences
|
||||
fa_features = parse_flash_attn_features()
|
||||
|
||||
# Parse FlashInfer TRTLLM feature differences (native vs TRTLLM on Blackwell)
|
||||
fi_features = parse_flashinfer_trtllm_features()
|
||||
|
||||
# Parse MLA prefill backends
|
||||
mla_prefill_backends = parse_mla_prefill_backends()
|
||||
|
||||
# Collect backend info
|
||||
all_backends = []
|
||||
for backend_name, class_path in attention_backends_map.items():
|
||||
if backend_name in SKIP_BACKENDS:
|
||||
@@ -1973,17 +1724,14 @@ def generate_docs() -> str:
|
||||
info = analyze_backend(backend_name, class_path)
|
||||
if info:
|
||||
all_backends.append(info)
|
||||
|
||||
# Expand backends into version variants
|
||||
if fa_features:
|
||||
all_backends = _expand_flash_attn_variants(all_backends, fa_features)
|
||||
if fi_features:
|
||||
all_backends = _expand_flashinfer_variants(all_backends, fi_features)
|
||||
|
||||
# DeepSeek V4 (*_DSV4) decode backends and MiniMax M3 sparse backends each
|
||||
# get their own subsection rather than mixing into the main MLA / standard
|
||||
# tables (the ROCm V4 backend isn't flagged is_mla by the AST heuristic, so
|
||||
# filter purely on the name).
|
||||
# DeepSeek V4 (*_DSV4) and MiniMax M3 sparse backends get their own tables
|
||||
# rather than mixing into the main MLA / standard tables (the ROCm V4 backend
|
||||
# isn't flagged is_mla by the AST heuristic, so filter purely on the name).
|
||||
def _is_v4(b: dict[str, Any]) -> bool:
|
||||
return b["name"].endswith("_DSV4")
|
||||
|
||||
@@ -1999,112 +1747,21 @@ def generate_docs() -> str:
|
||||
if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b)
|
||||
]
|
||||
|
||||
# Generate documentation
|
||||
script_path = "tools/pre_commit/generate_attention_backend_docs.py"
|
||||
doc_lines = [
|
||||
"# Attention Backend Feature Support",
|
||||
"",
|
||||
f"This document is auto-generated by `{script_path}`.",
|
||||
"It shows the feature support for each registered attention backend",
|
||||
"based on the checks in `AttentionBackend.validate_configuration()`.",
|
||||
"",
|
||||
"**Do not edit this file manually.** Run the following command to",
|
||||
"regenerate it:",
|
||||
"",
|
||||
"```bash",
|
||||
f"python {script_path}",
|
||||
"```",
|
||||
"",
|
||||
]
|
||||
|
||||
# Add usage documentation
|
||||
doc_lines.append(generate_usage_section())
|
||||
|
||||
# Add priority section
|
||||
doc_lines.append(generate_priority_section(priorities))
|
||||
|
||||
# Add legend and feature tables
|
||||
doc_lines.append(generate_legend())
|
||||
standard_title = "Standard Attention (MHA, MQA, GQA) Backends"
|
||||
doc_lines.append(
|
||||
generate_markdown_table(non_mla_backends, standard_title, is_mla_table=False)
|
||||
)
|
||||
# Add footnotes for version/variant distinctions (in table order)
|
||||
footnotes = []
|
||||
if fi_features:
|
||||
footnotes.append(
|
||||
"> **†** FlashInfer Native is the regular FlashInfer path. XQA is the "
|
||||
"SM90 decode path exposed through FlashInfer's TRTLLM decode API. "
|
||||
"trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen "
|
||||
"via `--attention-config.use_trtllm_attention=0`."
|
||||
)
|
||||
if fa_features:
|
||||
footnotes.append(
|
||||
"> **\\*** Specify the FlashAttention version via "
|
||||
"`--attention-config.flash_attn_version=2`, `3`, or `4`. "
|
||||
"Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), "
|
||||
"FA2 otherwise."
|
||||
)
|
||||
if footnotes:
|
||||
doc_lines.append("\n>\n".join(footnotes) + "\n")
|
||||
|
||||
# Add MiniMax M3 sparse section (separate category after standard GQA)
|
||||
if minimax_backends:
|
||||
doc_lines.append(generate_minimax_section(minimax_backends))
|
||||
|
||||
# Add MLA section with prefill and decode backends
|
||||
doc_lines.append(
|
||||
generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends)
|
||||
)
|
||||
|
||||
return "\n".join(doc_lines)
|
||||
mla_sm100_annotations = {"FLASHINFER_MLA_SPARSE": "**\\***"}
|
||||
return {
|
||||
"priority-standard": _priority_block(
|
||||
priorities, "standard_sm100", "standard_default"
|
||||
),
|
||||
"priority-mla": _priority_block(
|
||||
priorities, "mla_sm100", "mla_default", mla_sm100_annotations
|
||||
),
|
||||
"table-standard": _feature_table(non_mla_backends, is_mla=False),
|
||||
"table-minimax": _feature_table(minimax_backends, is_mla=False),
|
||||
"table-mla-prefill": _mla_prefill_table(mla_prefill_backends),
|
||||
"table-mla-decode": _feature_table(mla_backends, is_mla=True),
|
||||
"table-mla-v4-decode": _feature_table(v4_decode_backends, is_mla=True),
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate attention backend documentation table"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=str,
|
||||
default=str(REPO_ROOT / "docs" / "design" / "attention_backends.md"),
|
||||
help="Output file path (default: docs/design/attention_backends.md)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--check",
|
||||
action="store_true",
|
||||
help="Check if the documentation is up to date (for pre-commit)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"files",
|
||||
nargs="*",
|
||||
help="Files to check (passed by pre-commit). If none are relevant, skip.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.files and not any(is_relevant_file(f) for f in args.files):
|
||||
sys.exit(0)
|
||||
|
||||
output_path = Path(args.output)
|
||||
new_content = generate_docs()
|
||||
|
||||
if args.check:
|
||||
needs_update = (
|
||||
not output_path.exists() or output_path.read_text() != new_content
|
||||
)
|
||||
if needs_update:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(new_content)
|
||||
print(f"🔄 Regenerated: {output_path}")
|
||||
sys.exit(1)
|
||||
print(f"✅ Up to date: {output_path}")
|
||||
sys.exit(0)
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(new_content)
|
||||
print(f"Generated: {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
logger.info("Generating attention backend documentation")
|
||||
fill_markers("design/attention_backends.md", build_blocks())
|
||||
+34
-41
@@ -5,16 +5,15 @@ import logging
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
import mkdocs_awesome_nav.nav.directory as _nav_dir
|
||||
import mkdocs_gen_files
|
||||
import regex as re
|
||||
|
||||
logger = logging.getLogger("mkdocs")
|
||||
|
||||
ROOT_DIR = Path(__file__).parent.parent.parent.parent
|
||||
ROOT_DIR_RELATIVE = "../../../../.."
|
||||
EXAMPLE_DIR = ROOT_DIR / "examples"
|
||||
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/examples"
|
||||
|
||||
|
||||
def title(text: str) -> str:
|
||||
@@ -197,44 +196,38 @@ class Example:
|
||||
return content
|
||||
|
||||
|
||||
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
|
||||
# Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
|
||||
# title-cased (e.g. "Offline Inference" instead of "Offline inference").
|
||||
import mkdocs_awesome_nav.nav.directory as _nav_dir
|
||||
# Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
|
||||
# title-cased (e.g. "Offline Inference" instead of "Offline inference").
|
||||
_nav_dir.dirname_to_title = title
|
||||
logger.info("Generating example documentation")
|
||||
logger.debug("Root directory: %s", ROOT_DIR.resolve())
|
||||
logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
|
||||
|
||||
_nav_dir.dirname_to_title = title
|
||||
logger.info("Generating example documentation")
|
||||
logger.debug("Root directory: %s", ROOT_DIR.resolve())
|
||||
logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
|
||||
logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
|
||||
categories = sorted(
|
||||
p for p in EXAMPLE_DIR.iterdir() if p.is_dir() and not p.name.startswith(".")
|
||||
)
|
||||
|
||||
# Create the EXAMPLE_DOC_DIR if it doesn't exist
|
||||
if not EXAMPLE_DOC_DIR.exists():
|
||||
EXAMPLE_DOC_DIR.mkdir(parents=True)
|
||||
examples = []
|
||||
glob_patterns = ["*.py", "*.md", "*.sh"]
|
||||
# Find categorised examples
|
||||
for category in categories:
|
||||
logger.info("Processing category: %s", category.stem)
|
||||
globs = [category.glob(pattern) for pattern in glob_patterns]
|
||||
for path in itertools.chain(*globs):
|
||||
examples.append(Example(path, category.stem))
|
||||
# Find examples in subdirectories
|
||||
globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns]
|
||||
for path in itertools.chain(*globs):
|
||||
examples.append(Example(path.parent, category.stem))
|
||||
|
||||
categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
|
||||
|
||||
examples = []
|
||||
glob_patterns = ["*.py", "*.md", "*.sh"]
|
||||
# Find categorised examples
|
||||
for category in categories:
|
||||
logger.info("Processing category: %s", category.stem)
|
||||
globs = [category.glob(pattern) for pattern in glob_patterns]
|
||||
for path in itertools.chain(*globs):
|
||||
examples.append(Example(path, category.stem))
|
||||
# Find examples in subdirectories
|
||||
globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns]
|
||||
for path in itertools.chain(*globs):
|
||||
examples.append(Example(path.parent, category.stem))
|
||||
|
||||
# Generate the example documentation
|
||||
for example in sorted(examples, key=lambda e: e.path.stem):
|
||||
example_name = f"{example.path.stem}.md"
|
||||
doc_path = EXAMPLE_DOC_DIR / example.category / example_name
|
||||
if not doc_path.parent.exists():
|
||||
doc_path.parent.mkdir(parents=True)
|
||||
# Specify encoding for building on Windows
|
||||
with open(doc_path, "w+", encoding="utf-8") as f:
|
||||
f.write(example.generate())
|
||||
logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))
|
||||
logger.info("Total examples generated: %d", len(examples))
|
||||
# Generate the example documentation
|
||||
for example in sorted(examples, key=lambda e: e.path.stem):
|
||||
doc_path = f"examples/{example.category}/{example.path.stem}.md"
|
||||
with mkdocs_gen_files.open(doc_path, "w") as f:
|
||||
f.write(example.generate())
|
||||
if example.main_file is not None:
|
||||
# Point the edit button at the example's source file
|
||||
edit_path = Path("..") / example.main_file.relative_to(ROOT_DIR)
|
||||
mkdocs_gen_files.set_edit_path(doc_path, str(edit_path))
|
||||
logger.debug("Example generated: %s", doc_path)
|
||||
logger.info("Total examples generated: %d", len(examples))
|
||||
@@ -2,27 +2,28 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import ast
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from generated_content import fill_markers # noqa: E402
|
||||
|
||||
logger = logging.getLogger("mkdocs")
|
||||
|
||||
ROOT_DIR = Path(__file__).parent.parent.parent.parent
|
||||
DOCS_DIR = ROOT_DIR / "docs"
|
||||
GENERATED_METRICS_DIR = DOCS_DIR / "generated" / "metrics"
|
||||
|
||||
# Files to scan for metric definitions - each will generate a separate table
|
||||
# Files to scan for metric definitions - each fills a `gen:` marker in
|
||||
# docs/usage/metrics.md with its table (the section heading and any preamble
|
||||
# live in the tracked page next to the marker).
|
||||
METRIC_SOURCE_FILES = [
|
||||
{"path": "vllm/v1/metrics/loggers.py", "output": "general.inc.md"},
|
||||
{
|
||||
"path": "vllm/v1/spec_decode/metrics.py",
|
||||
"output": "spec_decode.inc.md",
|
||||
},
|
||||
{"path": "vllm/v1/metrics/loggers.py", "key": "metrics-general"},
|
||||
{"path": "vllm/v1/spec_decode/metrics.py", "key": "metrics-spec-decode"},
|
||||
{
|
||||
"path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py",
|
||||
"output": "nixl_connector.inc.md",
|
||||
"key": "metrics-nixl",
|
||||
},
|
||||
{"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"},
|
||||
{"path": "vllm/v1/metrics/perf.py", "key": "metrics-mfu"},
|
||||
]
|
||||
|
||||
|
||||
@@ -110,41 +111,27 @@ def generate_markdown_table(metrics: list[dict[str, str]]) -> str:
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
|
||||
"""Generate metrics documentation tables from source files."""
|
||||
logger.info("Generating metrics documentation")
|
||||
logger.info("Generating metrics documentation")
|
||||
|
||||
# Create generated directory if it doesn't exist
|
||||
GENERATED_METRICS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
blocks = {}
|
||||
total_metrics = 0
|
||||
for source_config in METRIC_SOURCE_FILES:
|
||||
source_path = source_config["path"]
|
||||
|
||||
total_metrics = 0
|
||||
for source_config in METRIC_SOURCE_FILES:
|
||||
source_path = source_config["path"]
|
||||
output_file = source_config["output"]
|
||||
filepath = ROOT_DIR / source_path
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"Metrics source file not found: {filepath}")
|
||||
|
||||
filepath = ROOT_DIR / source_path
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"Metrics source file not found: {filepath}")
|
||||
logger.debug("Extracting metrics from: %s", source_path)
|
||||
metrics = extract_metrics_from_file(filepath)
|
||||
logger.debug("Found %d metrics in %s", len(metrics), source_path)
|
||||
|
||||
logger.debug("Extracting metrics from: %s", source_path)
|
||||
metrics = extract_metrics_from_file(filepath)
|
||||
logger.debug("Found %d metrics in %s", len(metrics), source_path)
|
||||
blocks[source_config["key"]] = generate_markdown_table(metrics).strip()
|
||||
total_metrics += len(metrics)
|
||||
|
||||
# Generate and write the markdown table for this source
|
||||
table_content = generate_markdown_table(metrics)
|
||||
output_path = GENERATED_METRICS_DIR / output_file
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
f.write(table_content)
|
||||
|
||||
total_metrics += len(metrics)
|
||||
logger.info(
|
||||
"Generated metrics table: %s (%d metrics)",
|
||||
output_path.relative_to(ROOT_DIR),
|
||||
len(metrics),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Total metrics generated: %d across %d files",
|
||||
total_metrics,
|
||||
len(METRIC_SOURCE_FILES),
|
||||
)
|
||||
fill_markers("usage/metrics.md", blocks)
|
||||
logger.info(
|
||||
"Total metrics generated: %d across %d files",
|
||||
total_metrics,
|
||||
len(METRIC_SOURCE_FILES),
|
||||
)
|
||||
@@ -0,0 +1,56 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Inline build-time generated content into existing docs pages.
|
||||
|
||||
Source pages mark where generated content goes with a snippet-style marker,
|
||||
`--8<-- "gen:<key>"`, so the insertion point is explicit and readable. The
|
||||
substitution happens here (at gen-files time, before mkdocs-gen-files shadows
|
||||
the page), not via pymdownx.snippets, so the content can be generated at build
|
||||
time without living in a real file on disk.
|
||||
|
||||
The `gen:` prefix keeps these markers distinct from real pymdownx.snippets
|
||||
includes, and `fill_markers` fails loudly if a marker is missing or left behind
|
||||
(pymdownx.snippets would otherwise silently drop an unsubstituted marker).
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mkdocs_gen_files
|
||||
import regex as re
|
||||
|
||||
DOCS_DIR = Path(__file__).parent.parent.parent
|
||||
|
||||
_MARKER = '--8<-- "gen:{key}"'
|
||||
_ANY_MARKER = re.compile(r'--8<-- "gen:[^"]*"')
|
||||
|
||||
|
||||
def fill_markers(doc_path: str, blocks: dict[str, str]) -> None:
|
||||
"""Replace `--8<-- "gen:<key>"` markers in a docs page with generated content.
|
||||
|
||||
Args:
|
||||
doc_path: Docs-relative path of the source page to fill.
|
||||
blocks: Mapping of marker key to the markdown to insert in its place.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the source page does not exist.
|
||||
ValueError: If an expected marker is missing, or any `gen:` marker is
|
||||
left unsubstituted after filling.
|
||||
"""
|
||||
source = DOCS_DIR / doc_path
|
||||
if not source.exists():
|
||||
raise FileNotFoundError(f"Cannot fill markers in missing page: {doc_path}")
|
||||
|
||||
text = source.read_text()
|
||||
for key, content in blocks.items():
|
||||
marker = _MARKER.format(key=key)
|
||||
if marker not in text:
|
||||
raise ValueError(f"{doc_path}: missing marker {marker}")
|
||||
text = text.replace(marker, content)
|
||||
|
||||
if leftover := _ANY_MARKER.search(text):
|
||||
raise ValueError(f"{doc_path}: unsubstituted marker {leftover.group()}")
|
||||
|
||||
with mkdocs_gen_files.open(doc_path, "w") as f:
|
||||
f.write(text)
|
||||
# Keep the edit button pointing at the real source page
|
||||
mkdocs_gen_files.set_edit_path(doc_path, doc_path)
|
||||
@@ -19,6 +19,7 @@ The on_page_markdown hook passes the current page context to the preprocessor be
|
||||
each page is converted.
|
||||
"""
|
||||
|
||||
import posixpath
|
||||
from pathlib import Path
|
||||
|
||||
import regex as re
|
||||
@@ -38,18 +39,22 @@ TITLE = r"(?P<title>[^\[\]<>]+?)"
|
||||
REPO = r"(?P<repo>.+?/.+?)"
|
||||
TYPE = r"(?P<type>issues|pull|projects)"
|
||||
NUMBER = r"(?P<number>\d+)"
|
||||
VERSION = r"[^/\s]+"
|
||||
PATH = r"(?P<path>[^\s]+?)"
|
||||
FRAGMENT = r"(?P<fragment>#[^\s]+)?"
|
||||
URL = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}"
|
||||
URL_GITHUB = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}"
|
||||
RELATIVE = rf"(?!(https?|ftp)://|#){PATH}{FRAGMENT}"
|
||||
URL_DOCS = f"https://docs.vllm.ai/en/{VERSION}/{PATH}{FRAGMENT}"
|
||||
|
||||
# Common titles to use for GitHub links when none is provided in the link.
|
||||
TITLES = {"issues": "Issue ", "pull": "Pull Request ", "projects": "Project "}
|
||||
|
||||
# Regex to match GitHub issue, PR, and project links with optional titles.
|
||||
github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL}(\)|>)")
|
||||
github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL_GITHUB}(\)|>)")
|
||||
# Regex to match relative file links with optional titles.
|
||||
relative_link = re.compile(rf"\[{TITLE}\]\({RELATIVE}\)")
|
||||
# Regex to match absolute docs.vllm.ai links (should only exist in CLI).
|
||||
docs_link = re.compile(rf"\[{TITLE}\]\({URL_DOCS}\)")
|
||||
|
||||
|
||||
class UrlSchemesPreprocessor(Preprocessor):
|
||||
@@ -61,7 +66,8 @@ class UrlSchemesPreprocessor(Preprocessor):
|
||||
|
||||
def run(self, lines):
|
||||
page = self.ext.page
|
||||
if page is None or getattr(page.file, "abs_src_path", None) is None:
|
||||
files = self.ext.files
|
||||
if page is None:
|
||||
return lines
|
||||
|
||||
def replace_relative_link(match: re.Match) -> str:
|
||||
@@ -70,7 +76,7 @@ class UrlSchemesPreprocessor(Preprocessor):
|
||||
"""
|
||||
title = match.group("title")
|
||||
path = match.group("path")
|
||||
path = (Path(page.file.abs_src_path).parent / path).resolve()
|
||||
path = ((DOC_DIR / page.file.src_uri).parent / path).resolve()
|
||||
fragment = match.group("fragment") or ""
|
||||
|
||||
# Check if the path exists and is outside the docs dir
|
||||
@@ -105,9 +111,36 @@ class UrlSchemesPreprocessor(Preprocessor):
|
||||
url = f"https://github.com/{repo}/{type}/{number}{fragment}"
|
||||
return f"[{gh_icon} {title}]({url})"
|
||||
|
||||
def replace_docs_link(match: re.Match) -> str:
|
||||
"""Rewrite absolute docs.vllm.ai links as doc-relative links."""
|
||||
title = match.group("title")
|
||||
path = match.group("path").rstrip("/")
|
||||
fragment = match.group("fragment") or ""
|
||||
|
||||
# vllm.config.<Class> API reference -> mkdocstrings cross-reference
|
||||
if path == "api/vllm/config" and re.fullmatch(
|
||||
r"#vllm\.config\.\w+", fragment
|
||||
):
|
||||
ident = fragment[1:]
|
||||
return f"[`{ident}`][{ident}]"
|
||||
|
||||
# Other docs pages -> link relative to the current page, but only
|
||||
# when the target is a known docs page (real or generated); leave
|
||||
# unknown/external URLs untouched. This is correct even when the same
|
||||
# docstring is also rendered on its API reference page.
|
||||
src = f"{path.removesuffix('.html')}.md"
|
||||
if files.get_file_from_path(src) is None:
|
||||
return match.group(0)
|
||||
rel = posixpath.relpath(src, posixpath.dirname(page.file.src_uri))
|
||||
# Auto-wrapped bare URLs use the URL as their title; make it readable.
|
||||
if title.startswith("http"):
|
||||
title = path.removesuffix(".html")
|
||||
return f"[{title}]({rel}{fragment})"
|
||||
|
||||
markdown = "\n".join(lines)
|
||||
markdown = github_link.sub(replace_github_link, markdown)
|
||||
markdown = relative_link.sub(replace_relative_link, markdown)
|
||||
markdown = docs_link.sub(replace_docs_link, markdown)
|
||||
return markdown.split("\n")
|
||||
|
||||
|
||||
@@ -116,6 +149,7 @@ class UrlSchemesExtension(Extension):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.page = None
|
||||
self.files = None
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def extendMarkdown(self, md):
|
||||
@@ -138,4 +172,5 @@ def on_page_markdown(
|
||||
) -> str:
|
||||
"""Pass the current page context to the preprocessor."""
|
||||
_ext.page = page
|
||||
_ext.files = files
|
||||
return markdown
|
||||
|
||||
@@ -184,7 +184,7 @@ Our online Server provides endpoints that correspond to the offline APIs:
|
||||
- [Classification API](classify.md#online-serving)(`/classify`)
|
||||
- Corresponding to `LLM.score`:
|
||||
- [Score API](scoring.md#score-api) (`/score`, `/v1/score`)
|
||||
- [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
|
||||
- [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
|
||||
- Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models.
|
||||
|
||||
The following introduces the Pooling API. For other APIs, please refer to the link above.
|
||||
|
||||
@@ -20,7 +20,7 @@ The score models is designed to compute similarity scores between two input prom
|
||||
- `LLM.score`
|
||||
- Online APIs:
|
||||
- [Score API](scoring.md#score-api) (`/score`, `/v1/score`)
|
||||
- [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
|
||||
- [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
|
||||
|
||||
!!! note
|
||||
Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
|
||||
@@ -19,7 +19,7 @@ vLLM also supports model implementations that are available in Transformers. We
|
||||
|
||||
Currently, the Transformers modeling backend works for the following:
|
||||
|
||||
- Modalities: embedding models, language models and vision-language models*
|
||||
- Modalities: embedding models, language models, vision-language models* and audio-language models
|
||||
- Architectures: encoder-only, decoder-only, mixture-of-experts
|
||||
- Attention types: full attention and/or sliding attention
|
||||
|
||||
@@ -427,7 +427,6 @@ th {
|
||||
| `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
|
||||
| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |
|
||||
| `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ |
|
||||
| `OuroForCausalLM` | ouro | `ByteDance/Ouro-1.4B`, `ByteDance/Ouro-2.6B`, etc. | ✅︎ | |
|
||||
| `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |
|
||||
| `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
|
||||
| `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
|
||||
@@ -435,7 +434,6 @@ th {
|
||||
| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
|
||||
| `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ |
|
||||
| `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ |
|
||||
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -466,6 +464,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
|
||||
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
|
||||
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
|
||||
| `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
|
||||
@@ -608,7 +607,8 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
|
||||
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
| ------------ | ------ | ------ | ----------------- | --------------------------- | --------------------------------------- |
|
||||
| `Emu3ForConditionalGeneration` | Emu3 | T + I | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ |
|
||||
| `Emu3ForConditionalGeneration` | Emu3 | T + I<sup>+</sup> | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ |
|
||||
| `VibeVoiceAsrForConditionalGeneration` | VibeVoice-ASR | T + A<sup>+</sup> | `microsoft/VibeVoice-ASR-HF` | ✅︎ | ✅︎ |
|
||||
|
||||
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
|
||||
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
|
||||
|
||||
@@ -28,7 +28,7 @@ DOCS_PATHS=(
|
||||
docs/ # Actual docs content
|
||||
examples/ # Examples are rendered in docs
|
||||
vllm/ # API & CLI reference
|
||||
requirements/test/cuda.txt # CLI reference (see docs/mkdocs/hooks/generate_argparse.py)
|
||||
requirements/test/cuda.txt # CLI reference (see docs/mkdocs/gen_files/generate_argparse.py)
|
||||
mkdocs.yaml # Affects build process
|
||||
.readthedocs.yaml # Affects build process
|
||||
requirements/docs.txt # Affects build process
|
||||
|
||||
@@ -12,6 +12,11 @@ Before using EP, you need to install the necessary dependencies. We are actively
|
||||
2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation).
|
||||
3. **For disaggregated serving**: Install `gdrcopy` by running the [`install_gdrcopy.sh`](../../tools/install_gdrcopy.sh) script (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/).
|
||||
|
||||
!!! note "NCCL version (CUDA 13+)"
|
||||
The `deepep_v2` backend requires NCCL >= 2.30.4. PyTorch ships an older
|
||||
NCCL, so you must upgrade it before building or running DeepEP. See the
|
||||
[EP kernels guide](../../tools/ep_kernels) for instructions.
|
||||
|
||||
### Backend Selection Guide
|
||||
|
||||
vLLM provides multiple communication backends for EP. Use `--all2all-backend` to select one:
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user