Compare commits

..
1 Commits
Author SHA1 Message Date
+11 f68f4fddea kimi-k3
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Isotr0py <mozf@inferact.ai>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Summer Yang <girasoleyang@gmail.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: khluu <khluu000@gmail.com>
Co-authored-by: Bowen Wang <abmfy@icloud.com>
Co-authored-by: gnovack <novackgm@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com>
Co-authored-by: Roy Wang <yasong.wang@inferact.ai>
Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-07-27 07:01:52 +00:00
829 changed files with 11604 additions and 41823 deletions
-103
View File
@@ -1,103 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Audit vLLM compiled libraries for PyTorch stable ABI compliance."""
import fnmatch
import sys
from pathlib import Path
from torch_abi_audit import inspect_package
from torch_abi_audit.report import ExtensionReport, PackageReport
# Temporary allowlist of extensions not yet on the stable ABI.
# Shrink and remove over time.
ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
"_flashkda_C.abi3.so",
"vllm_flash_attn/_vllm_fa2_C.abi3.so",
"vllm_flash_attn/_vllm_fa3_C.abi3.so",
"third_party/deep_gemm/_C*.so",
)
def _relative_path(lib: ExtensionReport, package_root: Path) -> str:
try:
return lib.path.relative_to(package_root).as_posix()
except ValueError:
return lib.path.name
def _is_torch_unstable(lib: ExtensionReport) -> bool:
return lib.error is None and lib.torch.uses_torch and not lib.torch.stable
def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool:
return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns)
def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]:
return (*report.extensions, *report.bundled_libs)
def _collect_unstable(report: PackageReport) -> list[str]:
return sorted(
_relative_path(lib, report.root)
for lib in _iter_libs(report)
if _is_torch_unstable(lib)
)
def _find_stale_allowlist_entries(
report: PackageReport, patterns: tuple[str, ...]
) -> list[str]:
"""Allowlist patterns that match a built library which is no longer unstable."""
stale: list[str] = []
for pattern in patterns:
for lib in _iter_libs(report):
if lib.error is not None:
continue
if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern):
continue
if not _is_torch_unstable(lib):
stale.append(pattern)
break
return stale
def check_torch_abi(
package: str = "vllm",
patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES,
) -> int:
report = inspect_package(package)
if report.error:
print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr)
return 2
unstable = _collect_unstable(report)
unexpected = [
rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns)
]
stale = _find_stale_allowlist_entries(report, patterns)
if unexpected or stale:
if unexpected:
print(
"Not allowed: torch-unstable libraries outside "
f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}",
file=sys.stderr,
)
if stale:
print(
"Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: "
f"{', '.join(stale)}",
file=sys.stderr,
)
return 1
print("Torch stable ABI check passed.")
return 0
if __name__ == "__main__":
print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance")
sys.exit(check_torch_abi())
-1
View File
@@ -14,7 +14,6 @@ run_all_patterns:
- "setup.py"
- "csrc/"
- "cmake/"
- ".buildkite/check-torch-abi.py"
run_all_exclude_patterns:
- "docker/Dockerfile."
- "csrc/cpu/"
@@ -1,26 +0,0 @@
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/'
-76
View File
@@ -2,44 +2,6 @@ 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
@@ -61,41 +23,3 @@ 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"'
+4 -29
View File
@@ -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 (4 GPUs)
- label: NixlConnector PD accuracy (2 GPUs)
timeout_in_minutes: 60
num_devices: 4
num_devices: 2
device: intel_gpu
agent_tags:
label: production
gpu: 4+
gpu: 2+
mem: 16+
no_plugin: true
working_dir: "."
@@ -148,10 +148,7 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
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'
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
- label: Regression
key: regression
@@ -262,25 +259,3 @@ 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'
@@ -1,33 +0,0 @@
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: 24+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -28,9 +28,7 @@ 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)
@@ -62,55 +60,3 @@ 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'
-29
View File
@@ -1,29 +0,0 @@
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'
+2 -27
View File
@@ -145,32 +145,7 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_auto_round.py &&
pytest -v -s quantization/test_online.py'
- label: "XPU GPQA Eval (GPT-OSS)"
depends_on:
- image-build-xpu
timeout_in_minutes: 60
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/evals/gpt_oss/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install "gpt-oss[eval]==0.0.5" &&
cd tests &&
pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
pytest -v -s quantization/test_auto_round.py'
- label: "XPU compressed tensors FP8 test"
depends_on:
- image-build-xpu
@@ -193,4 +168,4 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
@@ -28,6 +28,11 @@
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
}
},
{
"dataset_name": "sharegpt",
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
}
},
{
"test_name": "serving_llama8B_tp1_random_128_128",
"server_parameters": {
-3
View File
@@ -7,9 +7,6 @@
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
+2 -19
View File
@@ -387,7 +387,6 @@ 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
@@ -406,14 +405,11 @@ initialize_native_environment() {
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_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
export PYTORCH_ROCM_ARCH=""
mkdir -p "${TMPDIR}" \
@@ -421,8 +417,7 @@ initialize_native_environment() {
"${TRITON_CACHE_DIR}" \
"${VLLM_CACHE_ROOT}" \
"${XDG_CACHE_HOME}" \
"${HF_HOME}" \
"${HF_DATASETS_CACHE}" || return 1
"${HF_HOME}" || return 1
echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
@@ -437,18 +432,6 @@ 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() {
@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
-e CMDS \
--name "${container_name}" \
"${IMAGE}" \
bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
>/dev/null
} 9>/tmp/docker-pull.lock
+40 -32
View File
@@ -169,6 +169,20 @@ 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
@@ -350,6 +364,22 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: V1 attention (H100-MI250) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
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
commands:
- pytest -v -s v1/attention
- label: V1 others (CPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -1547,12 +1577,11 @@ steps:
commands:
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: Kernels Core Operation Test %N # TBD
- label: Kernels Core Operation Test # 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/
@@ -1563,7 +1592,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 --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- 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
- label: Kernels KDA Test # TBD
timeout_in_minutes: 180
@@ -1576,29 +1605,10 @@ steps:
- vllm/third_party/flash_linear_attention/ops/kda.py
- vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
- vllm/third_party/flash_linear_attention/ops/l2norm.py
- vllm/models/kimi_k3/nvidia/kda.py
- vllm/models/kimi_k3/nvidia/kda_metadata.py
- vllm/models/kimi_k3/nvidia/ops/third_party/kda/
- tests/models/kimi_k3/test_kda.py
- tests/models/kimi_k3/test_kda_metadata.py
- tests/kernels/test_kda.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.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
- pytest -v -s kernels/test_kda.py
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
@@ -2151,7 +2161,7 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 8
parallelism: 4
optional: true
working_dir: "/vllm-workspace/"
source_file_dependencies:
@@ -2293,7 +2303,7 @@ steps:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
- pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
@@ -2684,12 +2694,11 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
- label: V1 attention (H100-MI300) %N # TBD
- label: V1 attention (H100-MI300) # 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:
@@ -2701,7 +2710,7 @@ steps:
- vllm/envs.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- pytest -v -s v1/attention
- label: V1 Core + KV + Metrics # TBD
timeout_in_minutes: 180
@@ -3667,12 +3676,11 @@ steps:
#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------#
- label: V1 attention (B200-MI355) %N # TBD
- label: V1 attention (B200-MI355) # 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
@@ -3683,7 +3691,7 @@ steps:
- vllm/envs.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- pytest -v -s v1/attention
- label: V1 Core + KV + Metrics # TBD
timeout_in_minutes: 180
-1
View File
@@ -16,7 +16,6 @@ steps:
commands:
- pytest -v -s cuda/test_cuda_context.py
- pytest -v -s cuda/test_platform_no_cuda_init.py
- pytest -v -s cuda/test_cuda_compatibility_path.py
- label: Cudagraph
device: h200_35gb
-16
View File
@@ -131,22 +131,6 @@ steps:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: NixlConnector PD edge case test (2 GPUs)
key: nixlconnector-pd-edge-cases-2-gpus
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- tests/v1/kv_connector/nixl_integration/
env:
PREFILL_GPU_ID: "0"
DECODE_GPU_ID: "1"
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
timeout_in_minutes: 25
-2
View File
@@ -40,11 +40,9 @@ steps:
source_file_dependencies:
- vllm/v1/engine/
- tests/v1/engine/
- tests/v1/test_tensor_ipc_queue.py
commands:
- pytest -v -s v1/engine/test_preprocess_error_handling.py
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
- pytest -v -s v1/test_tensor_ipc_queue.py
mirror:
amd:
device: mi250_1
@@ -52,5 +52,4 @@ steps:
- vllm/compilation/
- tests/distributed/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- pytest -v -s distributed/test_elastic_ep.py
@@ -1,26 +0,0 @@
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
+11 -36
View File
@@ -61,45 +61,9 @@ steps:
source_file_dependencies:
- csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
- vllm/models/deepseek_v4/common/ops/
- vllm/models/deepseek_v4/nvidia/
- tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
- tests/models/test_deepseek_v4_mega_moe.py
commands:
- pytest -v -s kernels/test_fused_deepseek_v4_*.py
- pytest -v -s models/test_deepseek_v4_mega_moe.py
# Catch-all for test files at the tests/kernels root. This job collects
# the whole root so new files are wired by default.
# Files with dedicated jobs elsewhere in this file are excluded via --ignore
# (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in
# their own jobs / Kernels (B200)).
- label: Kernels Root Misc Test (B200)
key: kernels-root-misc-test-b200
timeout_in_minutes: 45
device: b200-k8s
source_file_dependencies:
- csrc/
- vllm/
- tests/kernels/
commands:
- pytest -v -s kernels/
--ignore=kernels/attention
--ignore=kernels/core
--ignore=kernels/helion
--ignore=kernels/ir
--ignore=kernels/mamba
--ignore=kernels/moe
--ignore=kernels/quantization
--ignore=kernels/test_concat_mla_q.py
--ignore=kernels/test_fused_qk_norm_rope_gate.py
--ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
--ignore=kernels/test_top_k_per_row.py
--ignore=kernels/test_kda.py
--ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
--ignore=kernels/test_ll_bf16_gemm.py
--ignore=kernels/test_shuffle_rows.py
# BROKEN on main, pending kernel fixes (B200):
# test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)
- label: Kernels Attention Test %N
key: kernels-attention-test
@@ -214,6 +178,17 @@ steps:
commands:
- pytest -v -s kernels/mamba
- label: Kernels KDA Test
timeout_in_minutes: 25
device: h200_18gb
source_file_dependencies:
- vllm/third_party/flash_linear_attention/ops/kda.py
- vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
- vllm/third_party/flash_linear_attention/ops/l2norm.py
- tests/kernels/test_kda.py
commands:
- pytest -v -s kernels/test_kda.py
- label: Kernels DeepGEMM Test (H100)
key: kernels-deepgemm-test-h100
timeout_in_minutes: 35
+2 -2
View File
@@ -337,7 +337,7 @@ steps:
- label: LM Eval KV-Offload (2xH100)
key: kv-offload-medium
timeout_in_minutes: 45
timeout_in_minutes: 30
device: h100
num_devices: 2
source_file_dependencies:
@@ -347,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 or deepseek-v2-lite"
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
- label: LM Eval KV-Offload (4xH100)
key: kv-offload-large
+9 -16
View File
@@ -148,7 +148,6 @@ steps:
- 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 v1/test_kv_cache_spec_registry.py
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
@@ -266,7 +265,6 @@ steps:
- vllm/utils/
- vllm/v1/
- tests/v1/tracing
- tests/tracing/
commands:
- "pip install \
'opentelemetry-sdk>=1.26.0' \
@@ -274,7 +272,6 @@ steps:
'opentelemetry-exporter-otlp>=1.26.0' \
'opentelemetry-semantic-conventions-ai>=0.4.1'"
- pytest -v -s v1/tracing
- pytest -v -s tracing
mirror:
amd:
dind: false
@@ -298,7 +295,6 @@ steps:
amd:
device: mi250_1
timeout_in_minutes: 55
soft_fail: true
depends_on:
- image-build-amd
source_file_dependencies:
@@ -398,7 +394,7 @@ steps:
- label: Batch Invariance (A100)
key: batch-invariance-a100
timeout_in_minutes: 60
timeout_in_minutes: 40
device: a100
source_file_dependencies:
- vllm/v1/attention
@@ -408,11 +404,11 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- label: Batch Invariance (H100)
key: batch-invariance-h100
timeout_in_minutes: 60
timeout_in_minutes: 40
device: h100
source_file_dependencies:
- vllm/v1/attention
@@ -423,12 +419,12 @@ steps:
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- label: Batch Invariance (B200)
key: batch-invariance-b200
timeout_in_minutes: 45
timeout_in_minutes: 35
device: b200-k8s
source_file_dependencies:
- vllm/v1/attention
@@ -439,14 +435,11 @@ steps:
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
- pytest -v -s v1/determinism/test_matmul_batch_invariant.py
- pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
- pytest -v -s v1/determinism/test_online_batch_invariance.py
- label: Acceptance Length Test (Large Models) # optional
device: h200_35gb
key: acceptance-length-test-large-models
+1 -16
View File
@@ -61,20 +61,6 @@ steps:
# FA4 kernel tests require SM100; the suite skips them elsewhere.
- pytest -v -s models/inkling
- label: Kimi K3 Unit Tests (B200)
key: kimi-k3-unit-tests-b200
timeout_in_minutes: 40
device: b200-k8s
source_file_dependencies:
- vllm/models/kimi_k3/
- csrc/libtorch_stable/kimi_k3/
- tests/models/kimi_k3/
- tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
- tests/kernels/test_bf16_skinny_gemm.py
commands:
# The native NVIDIA Kimi K3 kernels require the SM100 family.
- pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
@@ -84,8 +70,7 @@ steps:
- vllm/
- tests/models/test_utils.py
- tests/models/test_vision.py
- tests/models/test_adapters.py
- tests/models/transformers/fusers/
device: cpu-small
commands:
- pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/
- pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
+1 -1
View File
@@ -26,7 +26,7 @@ steps:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
-14
View File
@@ -90,10 +90,8 @@ steps:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
- tests/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
- python3 spec_decode/test_custom_proposer.py
mirror:
amd:
dind: false
@@ -190,15 +188,3 @@ steps:
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
- label: Spec Decode Acceptance Rates Nightly
key: spec-decode-acceptance-rates-nightly
timeout_in_minutes: 60
device: h200_35gb
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
-14
View File
@@ -1,14 +0,0 @@
group: Torch ABI
depends_on:
- image-build
steps:
- label: Torch Stable ABI Audit
key: torch-stable-abi-audit
timeout_in_minutes: 5
source_file_dependencies:
- .buildkite/check-torch-abi.py
- csrc/
- cmake/
- setup.py
commands:
- python3 /vllm-workspace/.buildkite/check-torch-abi.py
+1
View File
@@ -3,6 +3,7 @@
dist
vllm/*.so
vllm/vllm-rs
.git
# Byte-compiled / optimized / DLL files
__pycache__/
-26
View File
@@ -19,7 +19,6 @@ pull_request_rules:
description: Comment on PR when pre-commit check fails
conditions:
- check-failure=pre-commit
- -check-cancelled=pre-commit
- -closed
- -draft
- or:
@@ -233,31 +232,6 @@ 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:
-19
View File
@@ -130,25 +130,6 @@ 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: [
{
+2 -2
View File
@@ -80,9 +80,9 @@ jobs:
'',
'\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.',
'',
'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.',
'PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.',
'',
'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.',
'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.',
'',
'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.',
'',
+1 -1
View File
@@ -41,7 +41,7 @@ jobs:
if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) {
core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`);
} else {
core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`);
core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`);
}
pre-commit:
-40
View File
@@ -1,40 +0,0 @@
name: Run CI from PR comment
on:
issue_comment:
types: [created]
concurrency:
group: run-ci-comment-${{ github.event.issue.number }}
cancel-in-progress: false
permissions:
contents: read
issues: write
pull-requests: write
jobs:
run-ci-command:
if: >-
github.event.issue.pull_request &&
(github.event.comment.body == '/ci run' ||
github.event.comment.body == '/ci retry')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
ref: ${{ github.event.repository.default_branch }}
persist-credentials: false
- uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
python-version: "3.12"
- name: Authorize and run CI command
run: >-
uv run --no-project --python 3.12
.github/workflows/scripts/run_ci_command.py
env:
BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
BUILDKITE_ORGANIZATION: vllm
BUILDKITE_PIPELINE: ci
CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }}
GH_TOKEN: ${{ github.token }}
-581
View File
@@ -1,581 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from collections.abc import Mapping, Sequence
from typing import Any
COMMAND_RUN_CI = "/ci run"
COMMAND_RETRY_FAILED = "/ci retry"
READY_LABELS = {"ready", "ready-run-all-tests"}
TRUSTED_PERMISSIONS = {"admin", "maintain", "write"}
ACTIVE_BUILD_STATES = {
"blocked",
"creating",
"scheduled",
"running",
"failing",
"canceling",
"waiting",
"waiting_failed",
}
RETRY_STATES = "failed,timed_out,expired"
class ApiError(RuntimeError):
def __init__(self, status: int | None, message: str) -> None:
super().__init__(message)
self.status = status
class HttpTransport:
def request(
self,
url: str,
*,
body: Mapping[str, Any] | None = None,
headers: Mapping[str, str] | None = None,
method: str = "GET",
) -> Any:
data = None if body is None else json.dumps(body).encode()
request = urllib.request.Request(
url,
data=data,
headers=dict(headers or {}),
method=method,
)
try:
with urllib.request.urlopen(request, timeout=30) as response:
response_body = response.read().decode()
except urllib.error.HTTPError as error:
response_body = error.read().decode()
message = self._error_message(response_body, error.reason)
raise ApiError(
error.code,
f"API returned {error.code}: {message}",
) from error
except urllib.error.URLError as error:
raise ApiError(None, f"API request failed: {error.reason}") from error
if not response_body:
return None
try:
return json.loads(response_body)
except json.JSONDecodeError as error:
raise ApiError(None, "API returned a non-JSON response.") from error
@staticmethod
def _error_message(response_body: str, fallback: str) -> str:
try:
parsed = json.loads(response_body)
except json.JSONDecodeError:
return fallback
return str(parsed.get("message", fallback))
class GitHubClient:
def __init__(
self,
token: str,
repository: str,
transport: HttpTransport | None = None,
) -> None:
if not token:
raise RuntimeError("GH_TOKEN is not set.")
self.owner, self.repo = repository.split("/", maxsplit=1)
self.transport = transport or HttpTransport()
self.headers = {
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"User-Agent": "vllm-ci-command",
"X-GitHub-Api-Version": "2022-11-28",
}
def _request(
self,
path: str,
*,
body: Mapping[str, Any] | None = None,
method: str = "GET",
) -> Any:
return self.transport.request(
f"https://api.github.com{path}",
body=body,
headers=self.headers,
method=method,
)
def _repo_path(self, suffix: str) -> str:
owner = urllib.parse.quote(self.owner, safe="")
repo = urllib.parse.quote(self.repo, safe="")
return f"/repos/{owner}/{repo}{suffix}"
def _paginate(self, path: str) -> list[dict[str, Any]]:
results: list[dict[str, Any]] = []
separator = "&" if "?" in path else "?"
for page in range(1, 101):
response = self._request(f"{path}{separator}per_page=100&page={page}")
if not isinstance(response, list):
raise ApiError(None, "GitHub API returned an invalid list response.")
results.extend(response)
if len(response) < 100:
return results
raise ApiError(None, "GitHub API pagination exceeded 10,000 results.")
def get_pr(self, number: int) -> dict[str, Any]:
return self._request(self._repo_path(f"/pulls/{number}"))
def get_permission(self, actor: str) -> str:
username = urllib.parse.quote(actor, safe="")
try:
response = self._request(
self._repo_path(f"/collaborators/{username}/permission")
)
except ApiError as error:
if error.status == 404:
return "none"
raise
return str(response["permission"])
def get_review_decision(self, number: int) -> str | None:
query = """
query($owner: String!, $repo: String!, $number: Int!) {
repository(owner: $owner, name: $repo) {
pullRequest(number: $number) {
reviewDecision
}
}
}
"""
response = self._request(
"/graphql",
body={
"query": query,
"variables": {
"number": number,
"owner": self.owner,
"repo": self.repo,
},
},
method="POST",
)
return response["data"]["repository"]["pullRequest"]["reviewDecision"]
def list_reviews(self, number: int) -> list[dict[str, Any]]:
return self._paginate(self._repo_path(f"/pulls/{number}/reviews"))
def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
return self._paginate(
self._repo_path(f"/issues/comments/{comment_id}/reactions")
)
def add_reaction(self, comment_id: int, content: str) -> None:
self._request(
self._repo_path(f"/issues/comments/{comment_id}/reactions"),
body={"content": content},
method="POST",
)
def add_comment(self, issue_number: int, body: str) -> None:
self._request(
self._repo_path(f"/issues/{issue_number}/comments"),
body={"body": body},
method="POST",
)
class BuildkiteClient:
def __init__(
self,
token: str,
organization: str,
pipeline: str,
transport: HttpTransport | None = None,
) -> None:
self.token = token
self.transport = transport or HttpTransport()
organization = urllib.parse.quote(organization, safe="")
pipeline = urllib.parse.quote(pipeline, safe="")
self.base_url = (
"https://api.buildkite.com/v2/organizations/"
f"{organization}/pipelines/{pipeline}/builds"
)
def _request(
self,
*,
body: Mapping[str, Any] | None = None,
method: str = "GET",
path: str = "",
query: Sequence[tuple[str, str]] = (),
) -> Any:
if not self.token:
raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.")
url = f"{self.base_url}{path}"
if query:
url = f"{url}?{urllib.parse.urlencode(query)}"
return self.transport.request(
url,
body=body,
headers={
"Authorization": f"Bearer {self.token}",
"Content-Type": "application/json",
"User-Agent": "vllm-ci-command",
},
method=method,
)
def list_builds(
self,
commit: str,
*,
metadata: tuple[str, str] | None = None,
) -> list[dict[str, Any]]:
query = [
("commit", commit),
("exclude_jobs", "true"),
("exclude_pipeline", "true"),
("per_page", "100"),
]
if metadata:
key, value = metadata
query.append((f"meta_data[{key}]", value))
response = self._request(query=query)
if not isinstance(response, list):
raise ApiError(None, "Buildkite API returned an invalid build list.")
return response
def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]:
return self._request(body=body, method="POST")
def retry_failed_jobs(
self,
build_number: int,
states: str,
) -> dict[str, Any]:
number = urllib.parse.quote(str(build_number), safe="")
return self._request(
body={"states": states},
method="PUT",
path=f"/{number}/retry_failed_jobs",
)
def parse_command(body: str) -> str | None:
if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}:
return body
return None
def parse_trusted_users(value: str = "") -> set[str]:
return {
user.casefold() for item in value.split(",") for user in item.split() if user
}
def has_ready_label(pr: Mapping[str, Any]) -> bool:
return any(label["name"] in READY_LABELS for label in pr["labels"])
def is_trusted_permission(permission: str) -> bool:
return permission in TRUSTED_PERMISSIONS
def authorize(
*,
actor: str,
permission: str,
pr: Mapping[str, Any],
trusted_approval: bool = False,
trusted_users: set[str] | None = None,
) -> tuple[bool, str]:
trusted_users = trusted_users or set()
if is_trusted_permission(permission):
return True, f"repository {permission} permission"
if actor.casefold() in trusted_users:
return True, "configured trusted contributor"
if actor.casefold() != pr["user"]["login"].casefold():
return (
False,
"Only reviewers with write access can run CI before it is "
"delegated to the PR author.",
)
if pr["draft"]:
return False, "PR authors cannot run CI while the PR is a draft."
if has_ready_label(pr):
return True, "ready label"
if trusted_approval:
return True, "approval from a trusted reviewer"
return (
False,
"A reviewer with write access must run `/ci run`, approve the PR, "
"or add the `ready` label first.",
)
def has_trusted_approval(
github: GitHubClient,
number: int,
trusted_users: set[str],
) -> bool:
if github.get_review_decision(number) != "APPROVED":
return False
latest_review_states: dict[str, tuple[str, str]] = {}
for review in github.list_reviews(number):
user = review.get("user") or {}
login = user.get("login")
state = review.get("state")
if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}:
latest_review_states[login.casefold()] = (login, state)
for login, state in latest_review_states.values():
if state != "APPROVED":
continue
if login.casefold() in trusted_users:
return True
if is_trusted_permission(github.get_permission(login)):
return True
return False
def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool:
pull_request = build.get("pull_request")
if isinstance(pull_request, Mapping):
build_pr_number = pull_request.get("id", pull_request.get("number"))
if build_pr_number is not None:
return str(build_pr_number) == str(pr_number)
metadata = build.get("meta_data") or {}
return str(metadata.get("github-pr-number")) == str(pr_number)
def is_active_build(build: Mapping[str, Any]) -> bool:
return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES
def select_latest_build(
builds: Sequence[dict[str, Any]],
pr_number: int,
) -> dict[str, Any] | None:
matching = [build for build in builds if is_build_for_pr(build, pr_number)]
return max(matching, key=lambda build: build.get("created_at", ""), default=None)
def create_build_payload(
*,
actor: str,
comment_id: int,
pr: Mapping[str, Any],
) -> dict[str, Any]:
return {
"commit": pr["head"]["sha"],
"branch": pr["head"]["ref"],
"message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}",
"pull_request_id": pr["number"],
"pull_request_base_branch": pr["base"]["ref"],
"pull_request_repository": pr["head"]["repo"]["clone_url"],
"pull_request_labels": [label["name"] for label in pr["labels"]],
"ignore_pipeline_branch_filters": True,
"env": {
"VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
"VLLM_CI_TRIGGERED_BY": actor,
},
"meta_data": {
"github-comment-id": str(comment_id),
"github-pr-number": str(pr["number"]),
"github-triggered-by": actor,
},
}
def add_reaction_safely(
github: GitHubClient,
comment_id: int,
content: str,
) -> None:
try:
github.add_reaction(comment_id, content)
except Exception as error:
print(f"Could not add {content} reaction: {error}", file=sys.stderr)
def is_already_handled(github: GitHubClient, comment_id: int) -> bool:
return any(
reaction.get("content") in {"rocket", "-1"}
and (reaction.get("user") or {}).get("login") == "github-actions[bot]"
for reaction in github.list_reactions(comment_id)
)
def handle_run_ci(
*,
actor: str,
buildkite: BuildkiteClient,
comment_id: int,
github: GitHubClient,
pr: Mapping[str, Any],
) -> str:
duplicate_builds = buildkite.list_builds(
pr["head"]["sha"],
metadata=("github-comment-id", str(comment_id)),
)
duplicate = select_latest_build(duplicate_builds, pr["number"])
if duplicate:
return f"CI was already requested by this comment: {duplicate['web_url']}"
current_builds = buildkite.list_builds(pr["head"]["sha"])
active_build = next(
(
build
for build in current_builds
if is_build_for_pr(build, pr["number"]) and is_active_build(build)
),
None,
)
if active_build:
return f"CI is already running for this commit: {active_build['web_url']}"
current_pr = github.get_pr(pr["number"])
if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
return (
"The PR head changed while processing the command. Comment `/ci run` again."
)
build = buildkite.create_build(
create_build_payload(
actor=actor,
comment_id=comment_id,
pr=current_pr,
)
)
return (
f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) "
f"for commit `{current_pr['head']['sha'][:12]}`."
)
def handle_retry_failed(
*,
buildkite: BuildkiteClient,
pr: Mapping[str, Any],
) -> str:
builds = buildkite.list_builds(pr["head"]["sha"])
build = select_latest_build(builds, pr["number"])
if not build:
return "No CI build exists for the current PR commit. Use `/ci run` first."
if not build.get("finished_at") or is_active_build(build):
return f"CI is still running for this commit: {build['web_url']}"
retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES)
if retried["retried_jobs_count"] == 0:
return (
f"No failed, timed-out, or expired jobs need retrying: {build['web_url']}"
)
return (
f"Queued {retried['retried_jobs_count']} failed job(s) for retry in "
f"[Buildkite CI #{build['number']}]({build['web_url']})."
)
def run(
event: Mapping[str, Any],
github: GitHubClient,
buildkite: BuildkiteClient,
trusted_users_value: str = "",
) -> None:
command = parse_command(event["comment"]["body"])
if not command or "pull_request" not in event["issue"]:
return
issue_number = event["issue"]["number"]
comment_id = event["comment"]["id"]
actor = event["comment"]["user"]["login"]
if is_already_handled(github, comment_id):
print(f"Comment {comment_id} was already handled.")
return
add_reaction_safely(github, comment_id, "eyes")
try:
pr = github.get_pr(issue_number)
permission = github.get_permission(actor)
if pr["state"] != "open":
github.add_comment(issue_number, "CI commands require an open PR.")
return
trusted_users = parse_trusted_users(trusted_users_value)
should_check_approval = (
not is_trusted_permission(permission)
and actor.casefold() not in trusted_users
and actor.casefold() == pr["user"]["login"].casefold()
and not pr["draft"]
and not has_ready_label(pr)
)
trusted_approval = should_check_approval and has_trusted_approval(
github,
issue_number,
trusted_users,
)
allowed, reason = authorize(
actor=actor,
permission=permission,
pr=pr,
trusted_approval=trusted_approval,
trusted_users=trusted_users,
)
if not allowed:
add_reaction_safely(github, comment_id, "-1")
github.add_comment(issue_number, f"@{actor}, {reason}")
return
print(f"Authorized @{actor}: {reason}")
if command == COMMAND_RUN_CI:
message = handle_run_ci(
actor=actor,
buildkite=buildkite,
comment_id=comment_id,
github=github,
pr=pr,
)
else:
message = handle_retry_failed(buildkite=buildkite, pr=pr)
add_reaction_safely(github, comment_id, "rocket")
github.add_comment(issue_number, message)
except Exception:
add_reaction_safely(github, comment_id, "confused")
raise
def main() -> None:
event_path = os.environ["GITHUB_EVENT_PATH"]
with open(event_path, encoding="utf-8") as event_file:
event = json.load(event_file)
if not parse_command(event["comment"]["body"]):
return
github = GitHubClient(
os.environ.get("GH_TOKEN", ""),
os.environ["GITHUB_REPOSITORY"],
)
buildkite = BuildkiteClient(
os.environ.get("BUILDKITE_API_TOKEN", ""),
os.environ.get("BUILDKITE_ORGANIZATION", "vllm"),
os.environ.get("BUILDKITE_PIPELINE", "ci"),
)
run(
event,
github,
buildkite,
os.environ.get("CI_TRUSTED_USERS", ""),
)
if __name__ == "__main__":
main()
@@ -1,363 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import unittest
from typing import Any
from run_ci_command import (
COMMAND_RETRY_FAILED,
COMMAND_RUN_CI,
RETRY_STATES,
BuildkiteClient,
authorize,
create_build_payload,
has_trusted_approval,
is_active_build,
is_build_for_pr,
parse_command,
parse_trusted_users,
run,
select_latest_build,
)
def make_pr(**overrides: Any) -> dict[str, Any]:
pr = {
"base": {"ref": "main"},
"draft": False,
"head": {
"ref": "feature",
"repo": {"clone_url": "https://github.com/contributor/vllm.git"},
"sha": "0123456789abcdef",
},
"labels": [],
"number": 42,
"state": "open",
"user": {"login": "author"},
}
pr.update(overrides)
return pr
def make_event(command: str, actor: str = "reviewer") -> dict[str, Any]:
return {
"comment": {
"body": command,
"id": 99,
"user": {"login": actor},
},
"issue": {
"number": 42,
"pull_request": {},
},
}
class FakeGitHub:
def __init__(
self,
*,
permission: str = "write",
permissions: dict[str, str] | None = None,
pr: dict[str, Any] | None = None,
review_decision: str = "REVIEW_REQUIRED",
reviews: list[dict[str, Any]] | None = None,
) -> None:
self.comments: list[str] = []
self.permission = permission
self.permissions = permissions or {}
self.pr = pr or make_pr()
self.reactions: list[str] = []
self.review_decision = review_decision
self.reviews = reviews or []
def get_pr(self, number: int) -> dict[str, Any]:
return self.pr
def get_permission(self, actor: str) -> str:
return self.permissions.get(actor, self.permission)
def get_review_decision(self, number: int) -> str:
return self.review_decision
def list_reviews(self, number: int) -> list[dict[str, Any]]:
return self.reviews
def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
return []
def add_reaction(self, comment_id: int, content: str) -> None:
self.reactions.append(content)
def add_comment(self, issue_number: int, body: str) -> None:
self.comments.append(body)
class FakeBuildkite:
def __init__(
self,
build_lists: list[list[dict[str, Any]]] | None = None,
) -> None:
self.build_lists = build_lists or []
self.created_builds: list[dict[str, Any]] = []
self.list_calls: list[tuple[str, tuple[str, str] | None]] = []
self.retry_calls: list[tuple[int, str]] = []
def list_builds(
self,
commit: str,
*,
metadata: tuple[str, str] | None = None,
) -> list[dict[str, Any]]:
self.list_calls.append((commit, metadata))
return self.build_lists.pop(0)
def create_build(self, body: dict[str, Any]) -> dict[str, Any]:
self.created_builds.append(body)
return {
"number": 123,
"web_url": "https://buildkite.example/builds/123",
}
def retry_failed_jobs(
self,
build_number: int,
states: str,
) -> dict[str, Any]:
self.retry_calls.append((build_number, states))
return {"retried_jobs_count": 3}
class FakeTransport:
def __init__(self, response: Any) -> None:
self.calls: list[dict[str, Any]] = []
self.response = response
def request(self, url: str, **kwargs: Any) -> Any:
self.calls.append({"url": url, **kwargs})
return self.response
class RunCiCommandTest(unittest.TestCase):
def test_only_exact_ci_commands_are_accepted(self) -> None:
self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI)
self.assertEqual(
parse_command(COMMAND_RETRY_FAILED),
COMMAND_RETRY_FAILED,
)
self.assertIsNone(parse_command("/ci run please"))
self.assertIsNone(parse_command(" /ci run"))
def test_write_access_authorizes_reviewers_and_authors(self) -> None:
allowed, _ = authorize(
actor="reviewer",
permission="write",
pr=make_pr(),
)
self.assertTrue(allowed)
def test_configured_trusted_contributors_can_run_ci(self) -> None:
trusted_users = parse_trusted_users("trusted-one, TRUSTED-TWO")
allowed, _ = authorize(
actor="trusted-two",
permission="read",
pr=make_pr(),
trusted_users=trusted_users,
)
self.assertTrue(allowed)
def test_authors_need_an_approval_or_ready_label(self) -> None:
pending, _ = authorize(
actor="author",
permission="read",
pr=make_pr(),
)
approved, _ = authorize(
actor="author",
permission="read",
pr=make_pr(),
trusted_approval=True,
)
ready, _ = authorize(
actor="author",
permission="read",
pr=make_pr(labels=[{"name": "ready"}]),
)
self.assertFalse(pending)
self.assertTrue(approved)
self.assertTrue(ready)
def test_non_author_contributors_without_write_are_denied(self) -> None:
allowed, _ = authorize(
actor="contributor",
permission="read",
pr=make_pr(),
trusted_approval=True,
)
self.assertFalse(allowed)
def test_authors_cannot_use_ready_state_on_draft_prs(self) -> None:
allowed, _ = authorize(
actor="author",
permission="read",
pr=make_pr(draft=True, labels=[{"name": "ready"}]),
trusted_approval=True,
)
self.assertFalse(allowed)
def test_only_trusted_reviewers_can_delegate_through_approval(self) -> None:
approved_review = {
"state": "APPROVED",
"user": {"login": "reviewer"},
}
trusted = FakeGitHub(
permission="read",
permissions={"reviewer": "write"},
review_decision="APPROVED",
reviews=[approved_review],
)
untrusted = FakeGitHub(
permission="read",
review_decision="APPROVED",
reviews=[approved_review],
)
self.assertTrue(has_trusted_approval(trusted, 42, set()))
self.assertFalse(has_trusted_approval(untrusted, 42, set()))
def test_build_matching_is_scoped_to_the_pr(self) -> None:
self.assertTrue(is_build_for_pr({"pull_request": {"id": 42}}, 42))
self.assertFalse(is_build_for_pr({"pull_request": {"id": 43}}, 42))
self.assertTrue(
is_build_for_pr(
{"meta_data": {"github-pr-number": "42"}},
42,
)
)
def test_latest_build_selection_ignores_other_prs(self) -> None:
latest = select_latest_build(
[
{
"created_at": "2026-07-28T02:00:00Z",
"number": 3,
"pull_request": {"id": 43},
},
{
"created_at": "2026-07-28T01:00:00Z",
"number": 2,
"pull_request": {"id": 42},
},
{
"created_at": "2026-07-28T00:00:00Z",
"number": 1,
"pull_request": {"id": 42},
},
],
42,
)
self.assertEqual(latest["number"], 2)
def test_active_build_states_prevent_duplicate_runs(self) -> None:
self.assertTrue(is_active_build({"state": "scheduled"}))
self.assertTrue(is_active_build({"state": "running"}))
self.assertTrue(is_active_build({"state": "waiting"}))
self.assertTrue(is_active_build({"blocked": True, "state": "passed"}))
self.assertFalse(is_active_build({"state": "failed"}))
def test_build_payload_preserves_pr_context(self) -> None:
payload = create_build_payload(
actor="reviewer",
comment_id=99,
pr=make_pr(labels=[{"name": "ready"}, {"name": "v1"}]),
)
self.assertEqual(
payload,
{
"commit": "0123456789abcdef",
"branch": "feature",
"message": "PR #42 /ci run by @reviewer",
"pull_request_id": 42,
"pull_request_base_branch": "main",
"pull_request_repository": ("https://github.com/contributor/vllm.git"),
"pull_request_labels": ["ready", "v1"],
"ignore_pipeline_branch_filters": True,
"env": {
"VLLM_CI_GITHUB_COMMENT_ID": "99",
"VLLM_CI_TRIGGERED_BY": "reviewer",
},
"meta_data": {
"github-comment-id": "99",
"github-pr-number": "42",
"github-triggered-by": "reviewer",
},
},
)
def test_ci_run_dispatches_build_with_current_pr_metadata(self) -> None:
github = FakeGitHub()
buildkite = FakeBuildkite([[], []])
run(make_event(COMMAND_RUN_CI), github, buildkite)
self.assertEqual(len(buildkite.created_builds), 1)
self.assertEqual(
buildkite.created_builds[0]["message"],
"PR #42 /ci run by @reviewer",
)
self.assertEqual(github.reactions, ["eyes", "rocket"])
self.assertIn("Buildkite CI #123", github.comments[0])
def test_unapproved_authors_are_denied_without_buildkite(self) -> None:
github = FakeGitHub(
permission="read",
pr=make_pr(),
review_decision="REVIEW_REQUIRED",
)
buildkite = FakeBuildkite()
run(make_event(COMMAND_RUN_CI, "author"), github, buildkite)
self.assertEqual(buildkite.list_calls, [])
self.assertEqual(github.reactions, ["eyes", "-1"])
self.assertIn("approve the PR", github.comments[0])
def test_ci_retry_uses_latest_current_sha_build(self) -> None:
github = FakeGitHub(
permission="read",
pr=make_pr(labels=[{"name": "ready"}]),
)
buildkite = FakeBuildkite(
[
[
{
"created_at": "2026-07-28T01:00:00Z",
"finished_at": "2026-07-28T02:00:00Z",
"number": 123,
"pull_request": {"id": 42},
"state": "failed",
"web_url": "https://buildkite.example/builds/123",
}
]
]
)
run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite)
self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)])
self.assertIn("Queued 3 failed job", github.comments[0])
def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None:
transport = FakeTransport({"retried_jobs_count": 2})
client = BuildkiteClient(
"secret",
"vllm",
"ci",
transport=transport,
)
client.retry_failed_jobs(123, RETRY_STATES)
call = transport.calls[0]
self.assertEqual(call["method"], "PUT")
self.assertTrue(call["url"].endswith("/123/retry_failed_jobs"))
self.assertEqual(call["body"], {"states": RETRY_STATES})
if __name__ == "__main__":
unittest.main()
+3
View File
@@ -173,6 +173,9 @@ venv.bak/
# mkdocs documentation
/site
docs/argparse
docs/examples/*
!docs/examples/README.md
# mypy
.mypy_cache/
-3
View File
@@ -3,9 +3,6 @@ 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
+4
View File
@@ -260,6 +260,10 @@ 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
+4 -9
View File
@@ -219,8 +219,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
# the set of architectures we want to compile for and remove the from the
# CMAKE_CUDA_FLAGS so that they are not applied globally.
#
# `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really
# needs PTX, add `+PTX` to that kernel's component-specific arch list below.
# `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch
# as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only
# `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's
# component-specific arch list below.
#
clear_cuda_arches(CUDA_ARCH_FLAGS)
extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}")
@@ -230,13 +232,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
cuda_archs_loose_intersection(CUDA_ARCHS
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
if(NOT CUDA_ARCHS)
message(FATAL_ERROR
"No supported CUDA architectures; the build would produce a binary "
"with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; "
"supported: ${CUDA_SUPPORTED_ARCHS}. "
"Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).")
endif()
else()
#
# For other GPU targets override the GPU architectures detected by cmake/torch
@@ -1358,10 +1358,6 @@ 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)
@@ -1,176 +0,0 @@
# 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())
@@ -186,7 +186,7 @@ def graph_samples(
outputs = [torch.empty_like(residual) for residual in residuals]
for weight, residual, output in zip(weights, residuals, outputs):
launch(activation, weight, residual, output)
torch.accelerator.synchronize()
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
@@ -194,7 +194,7 @@ def graph_samples(
launch(activation, weight, residual, output)
for _ in range(20):
graph.replay()
torch.accelerator.synchronize()
torch.cuda.synchronize()
samples = []
for _ in range(repeats):
@@ -276,7 +276,7 @@ def main() -> None:
raise ValueError("expected 1 <= M <= 16")
if not 0 <= args.config_shard < args.num_config_shards:
raise ValueError("config shard must be in [0, num_config_shards)")
torch.accelerator.set_device_index(0)
torch.cuda.set_device(0)
if torch.cuda.get_device_capability() != (10, 3):
raise RuntimeError("this benchmark requires SM103")
@@ -360,7 +360,7 @@ def main() -> None:
print(json.dumps(row, sort_keys=True), flush=True)
del activation, weights, residuals
torch.accelerator.empty_cache()
torch.cuda.empty_cache()
if __name__ == "__main__":
+1 -1
View File
@@ -154,7 +154,7 @@ def main(
scale=scale,
causal=True,
alibi_slopes=None,
sliding_window=window_size if sliding_window is not None else -1,
sliding_window=window_size,
block_table=block_tables,
softcap=0,
scheduler_metadata=metadata,
+1 -19
View File
@@ -15,7 +15,6 @@ 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")
@@ -97,14 +96,12 @@ 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
@@ -114,11 +111,6 @@ 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.
@@ -174,11 +166,6 @@ 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")
@@ -460,13 +447,8 @@ 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)
+2 -2
View File
@@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR)
message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
else()
# Keep in sync with tools/install_deepgemm.sh
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
set(_DEEPGEMM_UPSTREAM_REPO "git@github.com:Inferact/DeepGEMM.git")
# NOTE: This is currently targeting nv-dev branch due to sm120 support
set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _deepgemm_fc_root)
+1 -1
View File
@@ -12,7 +12,7 @@ if(FLASH_KDA_SRC_DIR)
else()
FetchContent_Declare(
flashkda
GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
GIT_REPOSITORY git@github.com:Inferact/FlashKDA.git
GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
GIT_PROGRESS TRUE
GIT_SUBMODULES cutlass
+5 -6
View File
@@ -241,15 +241,14 @@ endmacro()
# `<major>.<minor>`, dedupes them and then sorts them in ascending order and
# stores them in `OUT_ARCHES`.
#
# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags.
# This handles mismatches such as `arch=compute_20,code=sm_121`.
# Example:
# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a"
# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS)
# OUT_ARCHES="7.5;...;9.0"
function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS)
set(_CUDA_ARCHES)
foreach(_ARCH ${CUDA_ARCH_FLAGS})
string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
if (NOT _COMPUTE)
string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
endif()
string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
if (_COMPUTE)
set(_COMPUTE ${CMAKE_MATCH_1})
endif()
-11
View File
@@ -172,15 +172,4 @@
#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
+3 -6
View File
@@ -30,8 +30,7 @@ torch::Tensor get_scheduler_metadata(
const torch::Tensor& query_start_loc, const bool causal,
const int64_t window_size, const std::string& isa_hint,
const bool enable_kv_split,
const std::optional<torch::Tensor>& dynamic_causal,
const std::string& kv_cache_dtype) {
const std::optional<torch::Tensor>& dynamic_causal) {
cpu_attention::ISA isa;
if (isa_hint == "amx") {
isa = cpu_attention::ISA::AMX;
@@ -66,11 +65,9 @@ torch::Tensor get_scheduler_metadata(
input.dynamic_causal =
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
const int64_t kv_cache_idx =
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() {
input.elem_size = sizeof(attn_impl::kv_cache_t);
CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
input.elem_size = sizeof(scalar_t);
input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
input.output_buffer_elem_size =
+187 -5
View File
@@ -1,6 +1,5 @@
#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"
@@ -44,9 +43,193 @@
}()
namespace {
enum class FusedMOEAct {
SiluAndMul,
SwigluOAIAndMul,
GeluAndMul,
GeluTanhAndMul,
};
using cpu_fused_moe_utils::apply_gated_act;
using cpu_fused_moe_utils::FusedMOEAct;
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.");
}
}
template <typename scalar_t, typename gemm_t>
void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
@@ -634,7 +817,6 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
}
}
}
} // namespace
void prepack_moe_weight(
@@ -682,7 +864,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 = cpu_fused_moe_utils::get_act_type(act);
const FusedMOEAct act_type = 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");
-204
View File
@@ -1,204 +0,0 @@
// 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
-647
View File
@@ -1,647 +0,0 @@
// 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);
});
});
}
+3 -12
View File
@@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
explicit FP32Vec4(__vector float data) : reg(data) {}
FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
explicit 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) {}
FP32Vec8(const FP32Vec8& data) {
explicit 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) {}
FP32Vec16(const FP32Vec16& data) {
explicit 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,15 +747,6 @@ 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]);
+3 -3
View File
@@ -269,7 +269,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
explicit FP32Vec4(__vector float data) : reg(data) {}
FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
};
struct FP32Vec8 : public Vec<FP32Vec8> {
@@ -298,7 +298,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
explicit FP32Vec8(f32x4x2_t data) : reg(data) {}
FP32Vec8(const FP32Vec8& data) {
explicit FP32Vec8(const FP32Vec8& data) {
reg.val[0] = data.reg.val[0];
reg.val[1] = data.reg.val[1];
}
@@ -643,7 +643,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
FP32Vec16(const FP32Vec16& data) {
explicit 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];
@@ -31,9 +31,6 @@ 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,
@@ -1,424 +0,0 @@
// 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
+8 -14
View File
@@ -1,6 +1,3 @@
// 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
@@ -19,6 +16,9 @@ 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 += 4 * TileSize;
b_tile += Nr * K;
}
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 += 4 * TileSize;
b_tile1 += 4 * TileSize;
b_tile0 += Nr * K;
b_tile1 += Nr * K;
}
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
@@ -223,9 +223,6 @@ 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;
@@ -249,9 +246,6 @@ 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;
@@ -259,7 +253,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
public:
// physical layout [
// M / (8 or 4); Mr is 8 or 4
// M / 8; Mr is 8
// K / 4; K for bfmmla is 4
// 4, ; 4 row-pairs for each 8 rows
// 2, ; row-pair is 2 rows
@@ -445,7 +439,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 += NrGemv) {
for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
const bfloat16_t* __restrict__ b_panel =
reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
+12 -173
View File
@@ -451,90 +451,6 @@ 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]
@@ -629,7 +545,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, state_len) itype, where state_len >= width - 1
// conv_states: (..., dim, width - 1) itype
// activation: either None or "silu" or "swish"
// pad_slot_id: int
//
@@ -670,14 +586,11 @@ 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);
const int64_t state_len = conv_states_val.size(2);
CHECK_GE(state_len, width - 1);
CHECK_EQ(conv_states_val.size(2), 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);
@@ -738,14 +651,14 @@ at::Tensor causal_conv1d_fwd_cpu(
// API aligned with GPUs
//
// x: (batch, dim) or (batch, seqlen, dim)
// x: (batch, dim) or (batch, dim, seqlen)
// conv_state: (..., dim, state_len), where state_len >= width - 1
// weight: (dim, width)
// bias: (dim,)
// num_accepted_tokens: (batch,), dtype int32.
// cache_seqlens: (batch,), dtype int32.
// conv_state_indices: (batch,), dtype int32
// pad_slot_id: int
// out: (batch, dim) or (batch, seqlen, dim)
// out: (batch, dim) or (batch, dim, seqlen)
//
at::Tensor causal_conv1d_update_cpu(
const at::Tensor& x,
@@ -753,7 +666,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>& num_accepted_tokens,
const std::optional<at::Tensor>& cache_seqlens,
const std::optional<at::Tensor>& conv_state_indices,
int64_t pad_slot_id,
bool is_vnni) {
@@ -761,13 +674,13 @@ at::Tensor causal_conv1d_update_cpu(
CHECK_CONTIGUOUS(weight);
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
TORCH_CHECK(
x.dim() == 2 || x.dim() == 3,
"causal_conv1d_update_cpu: expect x to be 2D or 3D tensor.");
// 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.");
int64_t batch = x.size(0);
int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2);
int64_t seqlen = x.dim() == 2 ? 1 : x.size(1);
int64_t dim = x.size(1);
int64_t seqlen = 1;
int64_t width = weight.size(-1);
const auto scalar_type = x.scalar_type();
@@ -777,84 +690,10 @@ at::Tensor causal_conv1d_update_cpu(
CHECK_EQ(conv_states.scalar_type(), scalar_type);
CHECK_EQ(conv_states.size(1), dim);
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.");
CHECK_EQ(conv_states.size(2), width - 1);
// 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});
+7 -38
View File
@@ -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>& num_accepted_tokens,
bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
bool is_vnni);
@@ -163,8 +163,7 @@ torch::Tensor get_scheduler_metadata(
const torch::Tensor& query_start_loc, const bool casual,
const int64_t window_size, const std::string& isa_hint,
const bool enable_kv_split,
const std::optional<torch::Tensor>& dynamic_causal,
const std::string& kv_cache_dtype);
const std::optional<torch::Tensor>& dynamic_causal);
void cpu_attn_reshape_and_cache(const torch::Tensor& key,
const torch::Tensor& value,
@@ -208,20 +207,6 @@ 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,
@@ -517,8 +502,7 @@ 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? num_accepted_tokens, Tensor? conv_state_indices, int "
"pad_slot_id, "
"Tensor? cache_seqlens, 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
@@ -578,8 +562,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
"int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
"query_start_loc, bool casual, int window_size, str isa_hint, bool "
"enable_kv_split, Tensor? dynamic_causal, "
"str kv_cache_dtype=\"auto\") -> Tensor",
"enable_kv_split, Tensor? dynamic_causal) -> Tensor",
&get_scheduler_metadata);
ops.def(
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
@@ -613,7 +596,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#endif
// fused moe
#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
#if defined(__AVX512F__) || \
(defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
ops.def(
"prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
"-> ()");
@@ -624,22 +608,7 @@ 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 // #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__)
#endif
ops.def(
"mla_decode_kvcache("
" Tensor! out, Tensor query, Tensor kv_cache,"
+2 -3
View File
@@ -65,9 +65,8 @@ DINLINE LamportPack<P> load_lamport_pack(const P* ptr) {
: "l"(ptr)
: "memory");
#else
const volatile uint32_t* src =
reinterpret_cast<const volatile uint32_t*>(ptr);
#pragma unroll
const volatile uint32_t* src = reinterpret_cast<const volatile uint32_t*>(ptr);
#pragma unroll
for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
value.words[i] = src[i];
}
+1 -283
View File
@@ -3,155 +3,19 @@
#include <Python.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <filesystem>
#include <string>
#include <vector>
#if defined(O_DIRECT)
constexpr int kODirectFlag = O_DIRECT;
#else
constexpr int kODirectFlag = 0;
#endif
extern "C" {
namespace {
// Returns 0 on success, or the std::error_code's POSIX-compatible value on
// failure, mirroring the errno convention used by the syscalls below.
inline int ensure_parent_dirs(const std::string& path) {
const auto parent = std::filesystem::path(path).parent_path();
if (parent.empty()) {
return 0;
}
std::error_code ec;
std::filesystem::create_directories(parent, ec);
return ec ? ec.value() : 0;
}
// Core single-block store: src/size are raw pointer + byte count. Returns 0
// on success, or the errno of the failing step on failure -- captured
// before any subsequent cleanup call can overwrite it. On failure, the temp
// file is removed.
inline int _store_block(const char* tmp_path, const char* dest_path,
const char* src, size_t size, bool use_o_direct) {
if (access(dest_path, F_OK) == 0) {
return 0; // Already present.
}
if (const int err = ensure_parent_dirs(dest_path); err != 0) {
return err;
}
const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
const int fd = open(
tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644);
if (fd < 0) {
return errno;
}
const ssize_t written = write(fd, src, size);
if (written < 0 || static_cast<size_t>(written) != size) {
const int err = written < 0 ? errno : EIO;
close(fd); // Best-effort cleanup; the real error is already captured.
unlink(tmp_path);
return err;
}
if (close(fd) != 0) {
const int err = errno;
unlink(tmp_path);
return err;
}
if (rename(tmp_path, dest_path) != 0) {
const int err = errno;
unlink(tmp_path);
return err;
}
return 0;
}
// Core single-block load: dst/size are raw pointer + byte count. Returns 0
// on success, or the errno of the failing step on failure. On failure,
// the source file is removed since a partially-read block should not be reused.
inline int _load_block(const char* source_path, char* dst, size_t size,
bool use_o_direct) {
const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
const int fd = open(source_path, O_RDONLY | o_direct_flag, 0);
if (fd < 0) {
const int err = errno;
unlink(source_path);
return err;
}
const ssize_t bytes_read = read(fd, dst, size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != size) {
const int err = bytes_read < 0 ? errno : EIO;
close(fd);
unlink(source_path);
return err;
}
if (close(fd) != 0) {
const int err = errno;
unlink(source_path);
return err;
}
return 0;
}
inline void _batch_lookup(const std::vector<const char*>& paths,
static void _batch_lookup(const std::vector<const char*>& paths,
std::vector<int>& exists_flags) {
for (size_t i = 0; i < paths.size(); i++) {
exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0;
}
}
// Helper: extract a list[str] of length n into a vector<const char*>.
// Returns false and sets a Python exception on error.
inline bool extract_str_list(PyObject* list, Py_ssize_t n,
std::vector<const char*>& out) {
for (Py_ssize_t i = 0; i < n; i++) {
out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr);
if (out[i] == nullptr) {
return false;
}
}
return true;
}
// Helper: extract a Py_buffer per element of a list[bytes-like] of length n.
// On success, `out` holds n acquired buffers (caller must PyBuffer_Release
// each). On failure, any buffers already acquired are released before
// returning false, and a Python exception is set.
inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags,
std::vector<Py_buffer>& out) {
for (Py_ssize_t i = 0; i < n; i++) {
if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) {
for (Py_ssize_t j = 0; j < i; j++) {
PyBuffer_Release(&out[j]);
}
return false;
}
}
return true;
}
inline void release_buffer_list(std::vector<Py_buffer>& buffers) {
for (auto& buf : buffers) {
PyBuffer_Release(&buf);
}
}
} // namespace
/// @brief Check file existence for a batch of paths.
/// @param paths list[str] absolute paths to check.
/// @return list[bool] True if the corresponding path exists, False otherwise.
@@ -187,157 +51,11 @@ static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
return result;
}
/// @brief Store a batch of blocks, each from its own buffer, to disk.
/// @param tmp_paths list[str] one temp path per block.
/// @param dest_paths list[str] one destination path per block.
/// @param buffers list[bytes-like] one source buffer per block.
/// @param use_o_direct bool whether to open files with O_DIRECT
/// (default True). Ignored where O_DIRECT is unsupported
/// by the platform.
/// @note Releases the GIL for the entire batch. Raises on first error.
static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) {
PyObject* tmp_paths_obj = nullptr;
PyObject* dest_paths_obj = nullptr;
PyObject* buffers_obj = nullptr;
int use_o_direct = 1;
if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj,
&PyList_Type, &dest_paths_obj, &PyList_Type,
&buffers_obj, &use_o_direct)) {
return nullptr;
}
const Py_ssize_t n = PyList_Size(tmp_paths_obj);
if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) {
PyErr_SetString(
PyExc_ValueError,
"tmp_paths, dest_paths and buffers must have the same length");
return nullptr;
}
std::vector<const char*> tmp_paths(n);
std::vector<const char*> dest_paths(n);
if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr;
if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr;
std::vector<Py_buffer> buffers(n);
if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) {
return nullptr;
}
Py_ssize_t failed_index = -1;
int failure_errno = 0;
{
Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
const char* buf = static_cast<const char*>(buffers[i].buf);
const int err =
_store_block(tmp_paths[i], dest_paths[i], buf,
static_cast<size_t>(buffers[i].len), use_o_direct);
if (err != 0) {
failed_index = i;
failure_errno = err;
break;
}
}
Py_END_ALLOW_THREADS
}
release_buffer_list(buffers);
if (failed_index >= 0) {
// PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
errno = failure_errno;
return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
dest_paths[failed_index]);
}
Py_RETURN_NONE;
}
/// @brief Load a batch of blocks from disk, each into its own buffer.
/// @param source_paths list[str] one source path per block.
/// @param buffers list[writable bytes-like] one destination buffer
/// per block.
/// @param use_o_direct bool whether to open files with O_DIRECT
/// (default True). Ignored where O_DIRECT is unsupported
/// by the platform.
/// @note Releases the GIL for the entire batch. Raises on first error.
static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) {
PyObject* source_paths_obj = nullptr;
PyObject* buffers_obj = nullptr;
int use_o_direct = 1;
if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj,
&PyList_Type, &buffers_obj, &use_o_direct)) {
return nullptr;
}
const Py_ssize_t n = PyList_Size(source_paths_obj);
if (PyList_Size(buffers_obj) != n) {
PyErr_SetString(PyExc_ValueError,
"source_paths and buffers must have the same length");
return nullptr;
}
std::vector<const char*> source_paths(n);
if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr;
std::vector<Py_buffer> buffers(n);
if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) {
return nullptr;
}
Py_ssize_t failed_index = -1;
int failure_errno = 0;
{
Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
char* buf = static_cast<char*>(buffers[i].buf);
const int err =
_load_block(source_paths[i], buf, static_cast<size_t>(buffers[i].len),
use_o_direct);
if (err != 0) {
failed_index = i;
failure_errno = err;
break;
}
}
Py_END_ALLOW_THREADS
}
release_buffer_list(buffers);
if (failed_index >= 0) {
// PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
errno = failure_errno;
return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
source_paths[failed_index]);
}
Py_RETURN_NONE;
}
static PyMethodDef fs_io_C_methods[] = {
{"batch_lookup", batch_lookup, METH_VARARGS,
"batch_lookup(paths: list[str]) -> list[bool]\n"
"\n"
"Check file existence for a batch of paths."},
{"batch_store_block", batch_store_block, METH_VARARGS,
"batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n"
" buffers: list[bytes-like],\n"
" use_o_direct: bool = True) -> None\n"
"\n"
"Store a batch of blocks, each from its own buffer, to disk. Raises on "
"first error."},
{"batch_load_block", batch_load_block, METH_VARARGS,
"batch_load_block(source_paths: list[str],\n"
" buffers: list[writable bytes-like],\n"
" use_o_direct: bool = True) -> None\n"
"\n"
"Load a batch of blocks from disk into corresponding buffers. "
"Raises on first error."},
{nullptr, nullptr, 0, nullptr},
};
+12 -13
View File
@@ -697,8 +697,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
template <typename T, int kHdIn, int kHdOut, int kTileK = 256>
void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
int num_tokens, cudaStream_t const stream,
bool enable_pdl) {
int num_tokens, cudaStream_t const stream,
bool enable_pdl) {
if (num_tokens <= 8) {
invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK>(
output, mat_a, mat_b, num_tokens, stream, enable_pdl);
@@ -731,12 +731,10 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
// The kernels index global memory with raw pointers and packed strides, so
// reject any padded or transposed view rather than reading out of bounds.
STD_TORCH_CHECK(
mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
"mat_a must be a packed row-major [num_tokens, hd_in] tensor");
STD_TORCH_CHECK(
output.stride(0) == hd_out && output.stride(1) == 1,
"output must be a packed row-major [num_tokens, hd_out] tensor");
STD_TORCH_CHECK(mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
"mat_a must be a packed row-major [num_tokens, hd_in] tensor");
STD_TORCH_CHECK(output.stride(0) == hd_out && output.stride(1) == 1,
"output must be a packed row-major [num_tokens, hd_out] tensor");
STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in,
"mat_b must be a packed column-major [hd_in, hd_out] tensor");
@@ -760,11 +758,12 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
auto const* mat_b_ptr =
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \
if (hd_in == HD_IN && hd_out == HD_OUT) { \
invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \
output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \
return; \
#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \
if (hd_in == HD_IN && hd_out == HD_OUT) { \
invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \
output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, \
enable_pdl); \
return; \
}
// Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners
+4 -10
View File
@@ -249,9 +249,7 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
// For large num_tokens, use smaller blocks to increase SM concurrency.
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 grid(num_tokens);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -327,13 +325,8 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
/* This kernel is memory-latency bound in many scenarios.
When num_tokens is large, a smaller block size allows
for increased block occupancy on CUs and better latency
hiding on global mem ops. In batch-invariant mode the block size must
not depend on num_tokens, otherwise the same token would use a different
reduction width (and thus a different floating-point summation order)
across batches; lock it to 1024 to keep results bit-exact. */
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
hiding on global mem ops. */
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 block(std::min(hidden_size, max_block_size));
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -344,6 +337,7 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
bool offsets_are_multiple_of_vector_width =
hidden_size % vector_width == 0 && input_stride % vector_width == 0;
bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const bool has_weight = weight.has_value();
if (has_weight) {
auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());
@@ -215,9 +215,7 @@ void rms_norm_static_fp8_quant(
int num_tokens = input.numel() / hidden_size;
// For large num_tokens, use smaller blocks to increase SM concurrency.
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 grid(num_tokens);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -281,9 +279,7 @@ void fused_add_rms_norm_static_fp8_quant(
When num_tokens is large, a smaller block size allows
for increased block occupancy on CUs and better latency
hiding on global mem ops. */
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 block(std::min(hidden_size, max_block_size));
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -300,6 +296,7 @@ void fused_add_rms_norm_static_fp8_quant(
auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
bool ptrs_are_aligned =
inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&
!batch_invariant_launch) {
LAUNCH_FUSED_ADD_RMS_NORM(8);
@@ -2,7 +2,6 @@
#include "../../torch_utils.h"
#include "../../dispatch_utils.h"
#include "../../../core/batch_invariant.hpp"
#include "layernorm_utils.cuh"
#include "quant_conversions.cuh"
@@ -232,9 +231,7 @@ void rms_norm_per_block_quant_dispatch(
auto num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256);
const int max_block_size = (num_tokens <= 256) ? 512 : 256;
dim3 block(std::min(hidden_size, max_block_size));
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
+1 -1
View File
@@ -591,7 +591,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"limit=7.0) "
"-> ()");
// SituGLU implementation used in Kimi models.
// Kimi SITU (SituGLU) gated activation. linear_beta<=0 means unset.
ops.def(
"situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float "
"linear_beta=-1.0) -> ()");
+8 -30
View File
@@ -22,25 +22,17 @@ 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* d_flag_counters,
uint32_t data_offset, uint32_t flag_color,
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) \
@@ -50,21 +42,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, \
d_flag_counters, this->kMaxProblemSize); \
flag_color, 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, \
d_flag_counters, this->kMaxProblemSize); \
flag_color, 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, \
d_flag_counters, this->kMaxProblemSize); \
flag_color, this->kMaxProblemSize); \
}
// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
@@ -77,7 +69,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, \
d_flag_counters, this->kMaxProblemSize); \
flag_color, this->kMaxProblemSize); \
} else { \
throw std::runtime_error( \
"INT3 quick all-reduce is only supported for world_size == 2 " \
@@ -102,7 +94,7 @@ struct DeviceComms {
static int constexpr kMaxWorldSize = 8;
bool initialized = false;
uint32_t* d_flag_counters = nullptr;
uint32_t flag_color = 1;
int world_size;
int rank;
@@ -136,16 +128,6 @@ 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*)));
@@ -162,12 +144,6 @@ 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) {
@@ -235,6 +211,8 @@ struct DeviceComms {
break;
}
HIP_CHECK(cudaGetLastError());
// Rotate the flag color.
flag_color += divceil(N, grid);
}
};
+1 -1
View File
@@ -7,7 +7,7 @@ extern "C" {
#if defined(__i386__) || defined(__x86_64__)
#include <cpuid.h>
#include <x86intrin.h>
#include <mwaitxintrin.h>
#endif
#if defined(CLOCK_MONOTONIC_RAW)
+2 -2
View File
@@ -61,13 +61,13 @@ ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
FROM python-install AS torch-vision
# Install torchvision
ARG TORCH_VISION_VERSION=v0.28.0
ARG TORCH_VISION_VERSION=v0.26.0
WORKDIR /tmp
RUN --mount=type=cache,target=/root/.cache/uv \
git clone https://github.com/pytorch/vision.git && \
cd vision && \
git checkout $TORCH_VISION_VERSION && \
uv pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu && \
uv pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu && \
python setup.py bdist_wheel
FROM python-install AS hf-xet-builder
+35 -2
View File
@@ -63,6 +63,21 @@ RUN apt-get update -y && \
python3-pip && \
rm -rf /var/lib/apt/lists/*
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
printf '%s\n' \
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
'Pin: version 2025.3*' \
'Pin-Priority: 1001' \
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
apt-get update -y && \
apt-get install -y --no-install-recommends \
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
intel-oneapi-mkl-devel-2025.3 \
intel-oneapi-dnnl-devel-2025.3 && \
rm -rf /var/lib/apt/lists/*
# Install UMD
RUN mkdir neo && \
cd neo && \
@@ -85,6 +100,22 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
&& uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
rm "${ONECCL_INSTALLER}" && \
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
rm -f /opt/intel/oneapi/ccl/latest && \
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \
printf '%s\n' \
'/opt/intel/oneapi/ccl/2021.15/lib' \
'/opt/intel/oneapi/mpi/2021.15/lib' \
'/opt/intel/oneapi/compiler/2025.3/lib' \
> /etc/ld.so.conf.d/oneapi-ccl.conf && \
ldconfig
SHELL ["bash", "-c"]
CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
@@ -104,7 +135,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --upgrade pip
ENV LD_LIBRARY_PATH=/opt/venv/lib:/usr/local/lib
ENV LD_LIBRARY_PATH=/opt/intel/oneapi/ccl/2021.15/lib:/opt/intel/oneapi/mpi/2021.15/lib:/opt/intel/oneapi/compiler/2025.3/lib:/usr/local/lib
CMD ["/bin/bash"]
######################### UCX + NIXL BUILD STAGE #########################
@@ -184,7 +216,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r /workspace/vllm/requirements/xpu.txt && \
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
uv pip uninstall triton triton-xpu && \
uv pip install triton-xpu==3.7.2
uv pip install triton-xpu==3.7.1 && \
uv pip uninstall oneccl oneccl-devel
# Keep source-dependent layers near the end so frequent code-only changes
# don't invalidate heavy dependency and UCX/NIXL layers.
+1 -3
View File
@@ -56,9 +56,7 @@ nav:
- API Reference:
- api/README.md
- api/vllm
- CLI Reference:
- cli/README.md
- vllm: cli
- CLI Reference: cli
- Community:
- community/*
- Governance: governance
+9 -7
View File
@@ -1,8 +1,10 @@
nav:
- "*.md"
- bench:
- bench/*.md
- sweep:
- bench/sweep/*.md
- launch:
- launch/*.md
- README.md
- serve.md
- chat.md
- complete.md
- run-batch.md
- vllm bench:
- bench/**/*.md
- vllm launch:
- launch/**/*.md
+9
View File
@@ -0,0 +1,9 @@
# vllm bench latency
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_latency.inc.md"
+55
View File
@@ -0,0 +1,55 @@
# 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"
+9
View File
@@ -0,0 +1,9 @@
# vllm bench serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_serve.inc.md"
+9
View File
@@ -0,0 +1,9 @@
# vllm bench sweep plot
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_sweep_plot.inc.md"
+9
View File
@@ -0,0 +1,9 @@
# 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"
+9
View File
@@ -0,0 +1,9 @@
# vllm bench sweep serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_sweep_serve.inc.md"
+9
View File
@@ -0,0 +1,9 @@
# 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"
+9
View File
@@ -0,0 +1,9 @@
# vllm bench throughput
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_throughput.inc.md"
+5
View File
@@ -0,0 +1,5 @@
# vllm chat
## Arguments
--8<-- "docs/generated/argparse/chat.inc.md"
+5
View File
@@ -0,0 +1,5 @@
# vllm complete
## Arguments
--8<-- "docs/generated/argparse/complete.inc.md"
+10
View File
@@ -0,0 +1,10 @@
<!-- 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'`
+22
View File
@@ -0,0 +1,22 @@
# 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"
+9
View File
@@ -0,0 +1,9 @@
# vllm run-batch
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/run-batch.inc.md"
+9
View File
@@ -0,0 +1,9 @@
# vllm serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/serve.inc.md"
+9 -1
View File
@@ -11,4 +11,12 @@ 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<-- "gen:engine-args"
--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"
+2 -4
View File
@@ -301,10 +301,8 @@ review process:
isn't clear or you disagree with a suggestion, feel free to ask for
clarification or discuss the suggestion.
- Note that not all CI checks will be executed due to limited computational
resources. Reviewers with write access and configured trusted contributors
can comment `/ci run` when CI signals are needed before a PR is ready. After
the PR is approved or has the `ready` label, the PR author can use `/ci run`
or `/ci retry`. New commits do not start CI automatically.
resources. The reviewer will add `ready` label to the PR when the PR is
ready to merge or a full CI run is needed.
### Pull Request Limits and Escalation
+1 -1
View File
@@ -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 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.
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
- Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
Relevant server logic:
-20
View File
@@ -8,26 +8,6 @@ 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
+2 -34
View File
@@ -1,37 +1,5 @@
# llm-d
[llm-d](https://llm-d.ai/) is a Kubernetes-native distributed inference framework for serving large language models at scale, with vLLM as its primary inference engine. llm-d coordinates a fleet of vLLM instances across a cluster so that performance holds up under real production traffic, achieving the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators.
vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers.
It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d-to-the-cncf-evolving-kubernetes-into-sota-ai-infrastructure/) founded by Red Hat, Google Cloud, IBM Research, CoreWeave, and NVIDIA.
## What llm-d adds to vLLM
A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own:
- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it.
- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM.
- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts.
- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput.
- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch.
These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear.
## Performance
Representative benchmarked results across accelerators:
- **3x higher output throughput** and **2x faster TTFT** from prefix-aware routing vs round-robin (Llama 3.1 70B, AMD MI300X)
- **Up to 70% higher tokens/sec** from prefill/decode disaggregation (GPT-OSS, NVIDIA B200)
- **13.9x throughput** from hierarchical KV offloading at high concurrency vs GPU-only (NVIDIA H100)
See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and reproducible benchmarks on [Prism](https://prism.llm-d.ai/).
## Get started
1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration.
2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload.
3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment.
You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
Questions and contributions are welcome on [GitHub](https://github.com/llm-d/llm-d) and [Slack](https://llm-d.ai/slack).
You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
+90 -14
View File
@@ -1,9 +1,15 @@
# Attention Backend Feature Support
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()`.
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
```
## Setting the Attention Backend
@@ -92,11 +98,40 @@ Priority is **1 = highest** (tried first).
### Standard Attention (MHA, MQA, GQA)
--8<-- "gen:priority-standard"
**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` |
### MLA Attention (DeepSeek-style)
--8<-- "gen:priority-mla"
**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` |
> **\*** 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.
>
@@ -122,7 +157,24 @@ Priority is **1 = highest** (tried first).
## Standard Attention (MHA, MQA, GQA) Backends
--8<-- "gen:table-standard"
| 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 |
> **†** 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`.
>
@@ -136,7 +188,9 @@ 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.
--8<-- "gen:table-minimax"
| 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 |
## MLA (Multi-head Latent Attention) Backends
@@ -149,20 +203,38 @@ To explicitly select a prefill backend, use
Otherwise, the prefill backend is selected automatically at runtime based on
hardware and configuration.
--8<-- "gen:table-mla-prefill"
| 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 |
> **‡** Automatic selection tries FlashAttention first. On Blackwell
> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
> 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.
> 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`).
--8<-- "gen:table-mla-decode"
| 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 |
### DeepSeek V4 Decode Backends
@@ -173,4 +245,8 @@ 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.
--8<-- "gen:table-mla-v4-decode"
| 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 |
+1 -2
View File
@@ -101,7 +101,7 @@ When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks
Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
!!! note
Video CUDA graphs are automatically disabled when video token pruning (EVS or VidCom2) is enabled, since pruning makes the token count data-dependent and incompatible with CUDA graph capture.
Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
Mixed inputs (image+video) per prompt are also supported now.
@@ -129,7 +129,6 @@ 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` | ✅︎ | ❌︎ | ❌︎ |
-28
View File
@@ -49,34 +49,6 @@ Now supports 9 types of connectors:
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
```
## Reusing prefill token ids on decode
!!! note
This applies to disaggregated prefill and decode serving on the `/v1/chat/completions` endpoint, using a KV connector configured as in the Usage example above. It is experimental and subject to change.
In disaggregated serving, the prefill and decode stages both render the chat prompt from `messages` and tokenize it. Because the prefill stage has already produced the token ids, the decode stage can reuse them and skip its own templating and tokenization. The output is otherwise identical to a normal chat completion: it is detokenized to text, and tool and reasoning parsing, streaming, and structured output constraints all still apply.
The token ids are passed to the decode stage through `kv_transfer_params`, the dict already attached to the decode request to coordinate the transfer:
1. Send the prefill request with `return_token_ids` enabled, and read `prompt_token_ids` from the response.
2. Set `kv_transfer_params["prompt_token_ids"]` to those ids on the decode request. `messages` is still required, but its content is not tokenized when the ids are present.
```python
prefill = client.chat.completions.create(
model=model,
messages=messages,
extra_body={"return_token_ids": True, "kv_transfer_params": {"do_remote_decode": True}},
)
ids = prefill.prompt_token_ids
decode = client.chat.completions.create(
model=model,
messages=messages,
stream=True,
extra_body={"kv_transfer_params": {"do_remote_prefill": True, "prompt_token_ids": ids}},
)
```
## Development
We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.
+2 -98
View File
@@ -70,8 +70,7 @@ vllm serve <model> \
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: built-in `lru`/`arc`, or a custom `CachePolicy` name (see [Custom Eviction Policies](#custom-eviction-policies)). |
| `cache_policy_module_path` | no | — | both | Python import path for a custom `CachePolicy` not in the built-in registry. Required only when `eviction_policy` is not built-in and wasn't pre-registered via `CachePolicyFactory` (advanced). |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
@@ -79,36 +78,6 @@ vllm serve <model> \
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Custom Eviction Policies
`eviction_policy` resolves through `CachePolicyFactory` (`vllm/v1/kv_offload/cpu/policies/factory.py`), which pre-registers the built-in `lru` and `arc` policies.
### Out-of-tree (recommended)
Implement `CachePolicy` (`vllm/v1/kv_offload/cpu/policies/base.py`) in your own package — no vLLM fork or patch required — and point `kv_connector_extra_config` at it directly:
```json
{
"cpu_bytes_to_use": 10737418240,
"eviction_policy": "MyCachePolicy",
"cache_policy_module_path": "my_package.my_module"
}
```
`eviction_policy` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `cache_policy_module_path` and looks up `eviction_policy` as a class name in that module — the same fallback `spec_module_path` provides for a custom `OffloadingSpec`. No import or registration call needs to run before the server starts.
### Registering a friendly short name (in-process only)
If you control the process that constructs the vLLM engine (e.g. an embedding application), you can register a short name once at startup instead of repeating the module path in every config:
```python
from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
CachePolicyFactory.register_cache_policy("my_policy", "my_package.my_module", "MyCachePolicy")
```
Then set `"eviction_policy": "my_policy"` in `kv_connector_extra_config`, the same as `"lru"`/`"arc"`. This only takes effect within the process that ran the `register_cache_policy` call — it does not help when the server is launched as a separate process (e.g. via the `vllm serve` CLI), where the out-of-tree `cache_policy_module_path` config above is the only option.
## Secondary Tiers
Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
@@ -188,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.
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.
PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes.
| Key | Required | Default | Notes |
| --- | --- | --- | --- |
@@ -207,71 +176,6 @@ 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 24) 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.
+4 -46
View File
@@ -350,31 +350,6 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
#### Video Token Pruning
For supported models, vLLM can prune video tokens after the vision encoder to
reduce prefill time and KV cache usage, at some cost in accuracy. Set
`--video-pruning-rate <q>` to prune the fraction `q` of video tokens from each
video, and `--video-pruning-method` to choose the training-free algorithm:
- **`evs`** (Efficient Video Sampling, default): drops the tokens with the
lowest temporal dissimilarity to the previous frame. The first frame is
always fully retained.
- **`vidcom2`** (Video Compression Commander): scores tokens by similarity to
video-level and frame-level feature centers and gives distinctive frames a
larger share of the budget. At least one token per frame is retained.
```bash
vllm serve Qwen/Qwen3-VL-8B-Instruct \
--video-pruning-rate 0.75 --video-pruning-method vidcom2
```
!!! note
`evs` is supported by all models implementing multimodal pruning;
`vidcom2` is currently supported by Qwen3-VL only. Unsupported combinations
are rejected at startup. Enabling video pruning also disables encoder CUDA
graphs, since the retained token count becomes data-dependent.
### Audio Inputs
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
@@ -843,18 +818,16 @@ 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. Five
vLLM decodes video bytes into frames using a selectable decoding backend. Three
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.
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.
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.
Select the backend by passing the `backend` parameter via `--media-io-kwargs`:
@@ -881,21 +854,6 @@ 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.
-14
View File
@@ -19,20 +19,6 @@ 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.
@@ -42,12 +42,12 @@ pip install -v -r requirements/xpu.txt
```bash
pip uninstall -y triton triton-xpu
pip install triton-xpu==3.7.2 --extra-index-url https://download.pytorch.org/whl/xpu
pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu
```
!!! note
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
- For torch 2.13 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.2`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
- For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
- Finally, build and install vLLM XPU backend:
@@ -1,56 +0,0 @@
# 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)
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib.metadata
import importlib.util
import inspect
import logging
import sys
import textwrap
@@ -11,21 +10,17 @@ 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
from typing import TYPE_CHECKING, Literal
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):
@@ -137,8 +132,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"
@@ -147,28 +142,12 @@ 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")
@@ -200,7 +179,7 @@ class MarkdownFormatter(HelpFormatter):
def add_text(self, text: str):
if text:
self._markdown_output.append(f"{inspect.cleandoc(text)}\n\n")
self._markdown_output.append(f"{text.strip()}\n\n")
def add_usage(self, usage, actions, groups, prefix=None):
pass
@@ -262,163 +241,49 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser:
return _parser or parser
def format_help(parser: FlexibleArgumentParser) -> str:
"""Format a parser's help as markdown using `MarkdownFormatter`."""
return super(type(parser), parser).format_help()
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))
# 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))
if __name__ == "__main__":
on_startup("build", False)
@@ -5,15 +5,16 @@ 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:
@@ -196,38 +197,44 @@ class Example:
return content
# 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())
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
categories = sorted(
p for p in EXAMPLE_DIR.iterdir() if p.is_dir() and not p.name.startswith(".")
)
_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())
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))
# Create the EXAMPLE_DOC_DIR if it doesn't exist
if not EXAMPLE_DOC_DIR.exists():
EXAMPLE_DOC_DIR.mkdir(parents=True)
# 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))
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))

Some files were not shown because too many files have changed in this diff Show More