Compare commits

..
Author SHA1 Message Date
Bugen ZhaoandOpenAI Codex e1a763558c [CI] Discover Rust coverage artifacts from build metadata
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 06:44:16 +00:00
Bugen ZhaoandOpenAI Codex 84aeec9f22 [CI] Simplify Rust coverage reporting
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 08:18:14 +00:00
Bugen ZhaoandOpenAI Codex 82a770ddbd [CI] Simplify Rust coverage aggregation
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:58:45 +00:00
Bugen Zhao a09a9bace1 [CI] Disable redundant Codecov file fixes 2026-07-21 13:57:38 +00:00
Bugen Zhao 6c20d467a2 [CI] Run Codecov from repository root 2026-07-21 13:34:22 +00:00
Bugen ZhaoandOpenAI Codex cb59d0a351 [CI] Collect Rust coverage in Buildkite
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 13:01:56 +00:00
Bugen ZhaoandOpenAI Codex 0ab1bded36 [CI] Instrument Rust artifacts for coverage
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 12:22:25 +00:00
900 changed files with 11802 additions and 52539 deletions
-102
View File
@@ -1,102 +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, ...] = (
"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 -3
View File
@@ -145,8 +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'
pytest -v -s quantization/test_auto_round.py'
- label: "XPU compressed tensors FP8 test"
depends_on:
- image-build-xpu
@@ -169,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'
-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
+28 -28
View File
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_METADATA_VERSION="1"
IMAGE_EXISTED_BEFORE_BUILD=0
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
return 0
fi
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
if [[ -n "${!var:-}" ]]; then
echo "Using provided ${var}: ${!var}"
continue
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
compute_dependency_cache_keys() {
local bake_dir=""
local dockerfile_rocm=""
local nixl_branch=""
local rixl_branch=""
local ucx_branch=""
local rocshmem_branch=""
local deepep_branch=""
local nixl_material=""
local rixl_material=""
local rocshmem_material=""
local deepep_material=""
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
NIXL_CACHE_KEY=$(
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
RIXL_CACHE_KEY=$(
compose_dependency_cache_key \
"${nixl_branch}-ucx-${ucx_branch}" \
"${nixl_material}"
"${rixl_branch}-ucx-${ucx_branch}" \
"${rixl_material}"
)
export NIXL_CACHE_KEY
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
export RIXL_CACHE_KEY
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
fi
if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
case "${target}" in
nixl-rocm-ci)
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
rixl-rocm-ci)
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
fi
;;
rocshmem-rocm-ci)
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
resolve_ci_base_dependency_targets() {
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
local nixl_ref=""
local rixl_ref=""
local rocshmem_ref=""
local deepep_ref=""
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
case "${mode}" in
always)
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
add_dependency_cache_target "${target}"
fi
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
;;
esac
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
if dependency_cache_ref_exists "${nixl_ref}"; then
echo "NIXL dependency cache exists: ${nixl_ref}"
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
if dependency_cache_ref_exists "${rixl_ref}"; then
echo "RIXL dependency cache exists: ${rixl_ref}"
else
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
add_dependency_cache_target "nixl-rocm-ci"
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
add_dependency_cache_target "rixl-rocm-ci"
fi
fi
+9 -28
View File
@@ -35,7 +35,7 @@ set -o pipefail
: "${PY_COLORS:=1}"
: "${ROCM_DOCKER_TTY:=1}"
: "${PYTHONFAULTHANDLER:=1}"
: "${PYTEST_TIMEOUT:=2400}"
: "${PYTEST_TIMEOUT:=2100}"
if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
fi
@@ -45,9 +45,9 @@ fi
if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
fi
# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
# Dump stacks after 15 minutes, then stop an individual test after 35 minutes.
if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900"
fi
if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
" ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
@@ -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
@@ -401,19 +400,16 @@ initialize_native_environment() {
native_root="/tmp/vllm-native-${job_id}"
TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
VLLM_RPC_BASE_PATH="/tmp"
TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
TRITON_CACHE_DIR="${native_root}/cache/triton"
VLLM_CACHE_ROOT="${native_root}/cache/vllm"
XDG_CACHE_HOME="${native_root}/cache/xdg"
: "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}"
: "${TRITON_CACHE_DIR:=${native_root}/cache/triton}"
: "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}"
: "${XDG_CACHE_HOME:=${native_root}/cache/xdg}"
: "${HF_HOME:=/home/buildkite-agent/huggingface}"
# datasets uses POSIX locks that are unsupported by the shared HF NFS cache.
# Keep processed datasets job-local while retaining the persistent Hub cache.
HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets"
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
: "${HF_HUB_ETAG_TIMEOUT:=60}"
export TMPDIR VLLM_RPC_BASE_PATH
export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
export HF_HOME HF_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,10 +417,7 @@ initialize_native_environment() {
"${TRITON_CACHE_DIR}" \
"${VLLM_CACHE_ROOT}" \
"${XDG_CACHE_HOME}" \
"${HF_HOME}" \
"${HF_DATASETS_CACHE}" || return 1
echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
"${HF_HOME}" || return 1
if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
if ! command -v findmnt >/dev/null 2>&1; then
@@ -437,18 +430,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() {
@@ -1,11 +1,10 @@
#!/bin/bash
set -euox pipefail
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_CI_ENV=1
# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
# when callers specify fullgraph=True.
# Reduce sub-processes for acceleration
export TORCH_COMPILE_DISABLE=1
export VLLM_ENABLE_V1_MULTIPROCESSING=0
SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -50,15 +49,15 @@ wait_for_pid_and_check_log() {
}
# Test Sky Lake (AVX512F)
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
PID_TEST_0=$!
# Test Cascade Lake (AVX512F + VNNI)
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
PID_TEST_1=$!
# Test Cooper Lake (AVX512F + VNNI + BF16)
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
PID_TEST_2=$!
wait_for_pid_and_check_log $PID_TEST_0 test_0.log
@@ -29,6 +29,7 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
log_section() {
echo "--- $*"
@@ -106,6 +107,18 @@ install_cargo_nextest() {
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
}
install_cargo_llvm_cov() {
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
local toolchain
toolchain="$(rust_toolchain)"
rustup component add --toolchain "$toolchain" llvm-tools-preview
cargo binstall \
--no-confirm \
--force \
--secure \
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
}
install_uv() {
log_section "Installing uv ${UV_VERSION}"
curl -L --proto '=https' --tlsv1.2 -sSf \
@@ -176,14 +189,41 @@ run_tests() {
setup_pyo3_python
install_cargo_binstall
install_cargo_nextest
install_cargo_llvm_cov
log_section "Running cargo nextest"
cargo nextest run \
log_section "Running cargo nextest with Rust coverage"
mkdir -p artifacts
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
cargo llvm-cov clean \
--manifest-path rust/Cargo.toml \
--profraw-only
set +e
cargo llvm-cov nextest \
--manifest-path rust/Cargo.toml \
--workspace \
--all-features \
--locked \
--no-fail-fast
--no-fail-fast \
--no-clean \
--lcov \
--output-path artifacts/rust-unit.lcov \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
local coverage_rc=$?
local upload_rc=0
if [[ $coverage_rc -eq 0 ]]; then
# shellcheck source=.buildkite/scripts/rust-coverage.sh
source .buildkite/scripts/rust-coverage.sh
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
upload_rc=$?
fi
set -e
if [[ $coverage_rc -ne 0 ]]; then
return "$coverage_rc"
fi
return "$upload_rc"
}
install_protoc
+182
View File
@@ -0,0 +1,182 @@
#!/bin/sh
RUST_CODECOV_VERSION="v11.3.1"
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
rust_coverage_repo_root() {
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
printf '%s\n' /vllm-workspace
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
else
git rev-parse --show-toplevel
fi
}
rust_coverage_start() {
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
mkdir -p "$RUST_COVERAGE_DIR"
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
export LLVM_PROFILE_FILE
trap rust_coverage_finalize 0
}
rust_coverage_objects() {
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
python3 - "$rust_cov_objects_manifest" <<'PY'
from pathlib import Path
import sys
for relative in Path(sys.argv[1]).read_text().splitlines():
for entry in sys.path:
path = Path(entry or ".").resolve() / relative
if path.is_file():
print(path)
break
else:
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
PY
}
rust_coverage_collect() {
rust_cov_collect_flag=${1:?coverage flag is required}
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
rust_cov_collect_primary=
set --
while IFS= read -r rust_cov_collect_object; do
if [ -z "$rust_cov_collect_primary" ]; then
rust_cov_collect_primary=$rust_cov_collect_object
else
set -- "$@" "--object=$rust_cov_collect_object"
fi
done <<EOF
$rust_cov_collect_objects
EOF
llvm-profdata merge \
-sparse \
"$RUST_COVERAGE_DIR"/*.profraw \
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
llvm-cov export \
"$rust_cov_collect_primary" \
"$@" \
--format=lcov \
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
> "$rust_cov_collect_lcov" || return 1
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
export RUST_COVERAGE_LCOV
}
rust_coverage_upload() {
rust_cov_upload_lcov=${1:?LCOV path is required}
rust_cov_upload_flag=${2:?coverage flag is required}
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
if [ "$(uname -m)" != "x86_64" ]; then
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
return 1
fi
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|| return 1
curl -fsSL \
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
| sha256sum -c - || return 1
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
rust_cov_upload_slug="vllm-project/vllm"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
case "$rust_cov_upload_slug" in
*/*) ;;
*) rust_cov_upload_slug="vllm-project/vllm" ;;
esac
fi
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
if [ -z "${CODECOV_TOKEN:-}" ]; then
# Codecov accepts tokenless public uploads on unprotected branch names.
# A colon-separated prefix keeps feature-branch and fork uploads from
# requiring a repository secret.
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
else
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
fi
fi
set --
set -- "$@" upload-process
set -- "$@" --file "$rust_cov_upload_lcov"
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
# source-line fix scanning, which is unrelated to path mapping.
set -- "$@" --disable-search --disable-file-fixes
set -- "$@" --fail-on-error --git-service github
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
set -- "$@" --branch "$rust_cov_upload_branch"
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
set -- "$@" --slug "$rust_cov_upload_slug"
set -- "$@" --flag "$rust_cov_upload_flag"
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
set -- "$@" --dir "$rust_cov_upload_repo_root"
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
fi
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
# E2E steps run from tests/, so execute from the repository root to resolve
# codecov.yml and repository paths consistently.
(
cd "$rust_cov_upload_repo_root" || exit 1
"$rust_cov_upload_codecov_dir/codecov" "$@"
) >"$rust_cov_upload_log" 2>&1
rust_cov_upload_rc=$?
cat "$rust_cov_upload_log"
# v11.3.1 can log API failures while returning zero even with
# --fail-on-error. Preserve the strict CI contract explicitly.
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
echo "Codecov CLI reported an upload error" >&2
rust_cov_upload_rc=1
fi
rm -rf "$rust_cov_upload_codecov_dir"
return "$rust_cov_upload_rc"
}
rust_coverage_finalize() {
rust_cov_finalize_test_rc=$?
trap - 0
set +e
rust_coverage_collect "$RUST_COVERAGE_FLAG"
rust_cov_finalize_collect_rc=$?
rust_cov_finalize_upload_rc=0
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
rust_cov_finalize_upload_rc=$?
fi
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
exit "$rust_cov_finalize_test_rc"
fi
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
exit "$rust_cov_finalize_collect_rc"
fi
exit "$rust_cov_finalize_upload_rc"
}
+248 -292
View File
@@ -40,7 +40,7 @@
#####################################################################################################################################
# #
# IMPORTANT: #
# * Currently AMD CI has MI250 agents, MI300 agents, and MI355 agents. All upcoming feature improvements are #
# * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are #
# tracked in: https://github.com/vllm-project/vllm/issues/34994 #
# #
#-----------------------------------------------------------------------------------------------------------------------------------#
@@ -81,8 +81,10 @@
# the above test.) Also run if model initialization test file is modified. #
# * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model #
# source is modified, or when specified test files are modified. #
# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). #
# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). #
# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to #
# run plamo2 model in vLLM. #
# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) #
# and to run plamo2 model in vLLM. #
# * [Multi-Modal Models (Standard) 1-4]: #
# - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. #
# * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. #
@@ -169,6 +171,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
@@ -191,7 +207,6 @@ steps:
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -225,20 +240,6 @@ steps:
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
- label: Multi-Modal Processor (CPU) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
parallelism: 6
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
commands:
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------#
- label: Batch Invariance (H100-MI250) # TBD
@@ -350,21 +351,21 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: V1 others (CPU) # TBD
- label: V1 attention (H100-MI250) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/v1
- 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 -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
- pytest -v -s v1/attention
#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------#
@@ -407,19 +408,6 @@ steps:
- pytest -v -s transformers_utils
- pytest -v -s config
- label: Python-only Installation # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- tests/standalone_tests/python_only_compile.sh
- setup.py
- vllm/platforms/rocm.py
commands:
- bash standalone_tests/python_only_compile.sh
#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------#
- label: Rust Frontend Cargo Style + Clippy # TBD
@@ -457,7 +445,6 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
no_gpu: true
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- .buildkite/scripts/docker-build-metadata-args.sh
@@ -520,7 +507,7 @@ steps:
- tests/models/
commands:
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
- pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
@@ -670,30 +657,6 @@ steps:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/compilation/
- vllm/distributed/
- vllm/engine/
- vllm/executor/
- vllm/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- tests/compile/fullgraph/test_basic_correctness.py
- tests/compile/test_wrapper.py
- tests/entrypoints/llm/test_collective_rpc.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s entrypoints/llm/test_collective_rpc.py
- pytest -v -s ./compile/fullgraph/test_basic_correctness.py
- pytest -v -s ./compile/test_wrapper.py
#----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------#
- label: Platform Tests # TBD
@@ -716,7 +679,6 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -908,71 +870,6 @@ steps:
commands:
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/
- vllm/engine/
- vllm/executor/
- vllm/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- tests/distributed/
- tests/v1/shutdown
- tests/v1/worker/test_worker_memory_snapshot.py
- vllm/platforms/rocm.py
commands:
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
- label: Distributed Compile + Comm (4 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_4
num_gpus: 4
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/
- tests/distributed/test_pynccl
- tests/distributed/test_events
- tests/compile/fullgraph/test_basic_correctness.py
- tests/distributed/test_symm_mem_allreduce.py
- tests/distributed/test_multiproc_executor.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s compile/fullgraph/test_basic_correctness.py
- pytest -v -s distributed/test_pynccl.py
- pytest -v -s distributed/test_events.py
- pytest -v -s distributed/test_symm_mem_allreduce.py
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
#---------------------------------------------------------- mi300 · engine -----------------------------------------------------------#
- label: Engine # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/engine
- tests/test_sequence
- tests/test_config
- tests/test_logger
- tests/test_vllm_port
commands:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
- label: Entrypoints Unit Tests # TBD
@@ -981,7 +878,6 @@ steps:
dind: false
agent_pool: mi300_1
fast_check: true
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/entrypoints
@@ -1086,7 +982,6 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
fast_check: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -1115,7 +1010,6 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
fast_check: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -1130,7 +1024,6 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
fast_check: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -1452,27 +1345,6 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
- label: LM Eval Large Models (4xH100-4xMI300) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_USE_DEEP_GEMM=0
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------#
- label: Examples # TBD
@@ -1547,12 +1419,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 +1434,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
@@ -1581,21 +1452,6 @@ steps:
commands:
- pytest -v -s kernels/test_kda.py
- label: Kernels Mamba Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/mamba/
- tests/kernels/mamba
- vllm/model_executor/layers/mamba/ops
- vllm/platforms/rocm.py
commands:
- pytest -v -s kernels/mamba
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1737,7 +1593,7 @@ steps:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
- pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
- pytest -v -s v1/e2e/general/test_context_length.py
- pytest -v -s v1/e2e/general/test_min_tokens.py
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
@@ -1941,37 +1797,6 @@ steps:
- pip freeze | grep -E 'torch'
- pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- label: Language Models Test (Extended Generation) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/language/generation
commands:
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
- label: Language Models Tests (Hybrid) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/language/generation
commands:
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
#---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------#
- label: Multi-Modal Models (Extended Generation 1) # TBD
@@ -2074,32 +1899,20 @@ steps:
commands:
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Models (Extended Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal/pooling
commands:
- pytest -v -s models/multimodal/pooling -m 'not core_model'
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
- label: Multi-Modal Processor (CPU) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
commands:
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
#----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------#
@@ -2120,12 +1933,12 @@ steps:
#-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------#
- label: Transformers Nightly Models (Initialization) %N # TBD
- label: Transformers Nightly Models (Shardable) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 6
parallelism: 4
optional: true
working_dir: "/vllm-workspace/"
source_file_dependencies:
@@ -2141,27 +1954,6 @@ steps:
commands:
- pip install --upgrade git+https://github.com/huggingface/transformers
- pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- label: Transformers Nightly Models (Processing) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 8
optional: true
working_dir: "/vllm-workspace/"
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/multimodal/
- vllm/model_executor/layers/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- tests/models/
commands:
- pip install --upgrade git+https://github.com/huggingface/transformers
- pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- label: Transformers Nightly Models (Single) # TBD
@@ -2680,12 +2472,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:
@@ -2697,7 +2488,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
@@ -2726,6 +2517,23 @@ steps:
# - export HSA_NO_SCRATCH_RECLAIM=1
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
- label: V1 others (CPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/v1
commands:
- pytest -v -s -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
- label: V1 Sample + Logits # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -3022,6 +2830,195 @@ steps:
commands:
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
#########################################################################################################################################
# #
# MI325 (gfx942) tests #
# #
#########################################################################################################################################
#---------------------------------------------------------- mi325 · compile ----------------------------------------------------------#
- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/compilation/
- vllm/distributed/
- vllm/engine/
- vllm/executor/
- vllm/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- tests/compile/fullgraph/test_basic_correctness.py
- tests/compile/test_wrapper.py
- tests/entrypoints/llm/test_collective_rpc.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s entrypoints/llm/test_collective_rpc.py
- pytest -v -s ./compile/fullgraph/test_basic_correctness.py
- pytest -v -s ./compile/test_wrapper.py
#-------------------------------------------------------- mi325 · distributed --------------------------------------------------------#
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/
- vllm/engine/
- vllm/executor/
- vllm/worker/worker_base.py
- vllm/v1/engine/
- vllm/v1/worker/
- tests/distributed/
- tests/v1/shutdown
- tests/v1/worker/test_worker_memory_snapshot.py
- vllm/platforms/rocm.py
commands:
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
- label: Distributed Compile + Comm (4 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_4
num_gpus: 4
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/
- tests/distributed/test_pynccl
- tests/distributed/test_events
- tests/compile/fullgraph/test_basic_correctness.py
- tests/distributed/test_symm_mem_allreduce.py
- tests/distributed/test_multiproc_executor.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s compile/fullgraph/test_basic_correctness.py
- pytest -v -s distributed/test_pynccl.py
- pytest -v -s distributed/test_events.py
- pytest -v -s distributed/test_symm_mem_allreduce.py
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
#---------------------------------------------------------- mi325 · engine -----------------------------------------------------------#
- label: Engine # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/engine
- tests/test_sequence
- tests/test_config
- tests/test_logger
- tests/test_vllm_port
commands:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------#
- label: LM Eval Large Models (4xH100-4xMI325) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_USE_DEEP_GEMM=0
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
#----------------------------------------------------- mi325 · models / language -----------------------------------------------------#
- label: Language Models Test (Extended Generation) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/language/generation
commands:
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
- label: Language Models Tests (Hybrid) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_1
parallelism: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/language/generation
commands:
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
#---------------------------------------------------- mi325 · models / multimodal ----------------------------------------------------#
- label: Multi-Modal Models (Extended Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal/pooling
commands:
- pytest -v -s models/multimodal/pooling -m 'not core_model'
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal
commands:
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
#----------------------------------------------------------- mi325 · misc ------------------------------------------------------------#
- label: Python-only Installation # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- tests/standalone_tests/python_only_compile.sh
- setup.py
- vllm/platforms/rocm.py
commands:
- bash standalone_tests/python_only_compile.sh
#########################################################################################################################################
# #
# MI355 (gfx950) tests #
@@ -3033,7 +3030,6 @@ steps:
- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/"
@@ -3050,7 +3046,6 @@ steps:
- label: Distributed Tests (2xH100-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3095,7 +3090,6 @@ steps:
- label: Entrypoints Integration (API Server) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
fast_check: true
@@ -3113,7 +3107,6 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
optional: true
@@ -3129,7 +3122,6 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
optional: true
@@ -3146,7 +3138,6 @@ steps:
- label: Entrypoints Integration (API Server Generate) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
optional: true
@@ -3167,7 +3158,6 @@ steps:
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
working_dir: "/vllm-workspace/tests"
@@ -3181,7 +3171,6 @@ steps:
- label: Entrypoints Integration (Multimodal)
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
working_dir: "/vllm-workspace/tests"
@@ -3195,7 +3184,6 @@ steps:
- label: Entrypoints Integration (Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
working_dir: "/vllm-workspace/tests"
@@ -3211,7 +3199,6 @@ steps:
- label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3234,7 +3221,6 @@ steps:
- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3257,7 +3243,6 @@ steps:
- label: LM Eval Small Models (2xB200-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3277,7 +3262,6 @@ steps:
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace"
@@ -3298,7 +3282,6 @@ steps:
- label: LM Eval Large Models (4xH100-4xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_4
num_gpus: 4
optional: true
@@ -3321,7 +3304,6 @@ steps:
- label: Examples # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
@@ -3357,7 +3339,6 @@ steps:
- label: Kernels (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/"
source_file_dependencies:
@@ -3383,7 +3364,6 @@ steps:
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 2
working_dir: "/vllm-workspace/tests"
@@ -3401,7 +3381,6 @@ steps:
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 5
working_dir: "/vllm-workspace/tests"
@@ -3422,7 +3401,6 @@ steps:
- label: Kernels Quantization Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 2
working_dir: "/vllm-workspace/tests"
@@ -3440,7 +3418,6 @@ steps:
- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
@@ -3460,7 +3437,6 @@ steps:
- label: Language Models Test (Extended Generation) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3474,7 +3450,6 @@ steps:
- label: Language Models Test (Extended Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3487,9 +3462,7 @@ steps:
- label: Language Models Test (PPL) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/model_executor/models/qwen3_5.py
@@ -3516,7 +3489,6 @@ steps:
- label: Language Models Tests (Standard) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3531,7 +3503,6 @@ steps:
- label: Multi-Modal Models (Extended Generation 1) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3546,7 +3517,6 @@ steps:
- label: Multi-Modal Models (Extended Generation 3) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3559,7 +3529,6 @@ steps:
- label: Multi-Modal Models (Extended Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3572,7 +3541,6 @@ steps:
- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3586,7 +3554,6 @@ steps:
- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3603,7 +3570,6 @@ steps:
- label: Quantized Models Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3620,7 +3586,6 @@ steps:
- label: Quantization # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3637,7 +3602,6 @@ steps:
# - label: Quantized MoE Test (B200-MI355) # TBD
# timeout_in_minutes: 180
# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
# dind: false
# agent_pool: mi355_1
# working_dir: "/vllm-workspace/"
# source_file_dependencies:
@@ -3663,12 +3627,10 @@ 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
@@ -3679,12 +3641,11 @@ 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
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3711,7 +3672,6 @@ steps:
- label: V1 Sample + Logits # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3732,7 +3692,6 @@ steps:
- label: V1 Spec Decode # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3746,7 +3705,6 @@ steps:
- label: Weight Loading Multiple GPU # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
@@ -3759,7 +3717,6 @@ steps:
- label: Weight Loading Multiple GPU - Large Models # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
working_dir: "/vllm-workspace/tests"
num_gpus: 2
@@ -3775,7 +3732,6 @@ steps:
- label: Regression # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
+2 -3
View File
@@ -16,9 +16,8 @@ steps:
parallelism: 2
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 125
device: mi325_1
timeout_in_minutes: 95
depends_on:
- image-build-amd
source_file_dependencies:
+3 -4
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -18,8 +18,7 @@ steps:
- pytest -v -s basic_correctness/test_cpu_offload.py
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 60
device: mi325_1
timeout_in_minutes: 70
depends_on:
- image-build-amd
+1 -2
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 45
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -15,7 +15,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 40
depends_on:
- image-build-amd
-4
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
@@ -27,10 +26,7 @@ steps:
- vllm/v1/cudagraph_dispatcher.py
- vllm/config/compilation.py
- vllm/compilation
- vllm/v1/worker/encoder_cudagraph.py
- vllm/v1/worker/encoder_cudagraph_defs.py
commands:
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+5 -21
View File
@@ -17,7 +17,7 @@ steps:
amd:
dind: false
device: mi300_4
timeout_in_minutes: 60
timeout_in_minutes: 85
depends_on:
- image-build-amd
source_file_dependencies:
@@ -68,7 +68,7 @@ steps:
amd:
dind: false
device: mi300_4
timeout_in_minutes: 40
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
@@ -94,7 +94,7 @@ steps:
amd:
dind: false
device: mi300_4
timeout_in_minutes: 60
timeout_in_minutes: 85
depends_on:
- image-build-amd
source_file_dependencies:
@@ -120,7 +120,7 @@ steps:
amd:
dind: false
device: mi300_4
timeout_in_minutes: 55
timeout_in_minutes: 80
depends_on:
- image-build-amd
source_file_dependencies:
@@ -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
@@ -193,7 +177,7 @@ steps:
amd:
dind: false
device: mi300_2
timeout_in_minutes: 45
timeout_in_minutes: 70
depends_on:
- image-build-amd
source_file_dependencies:
-1
View File
@@ -41,7 +41,6 @@ steps:
amd:
dind: false
device: mi300_2
timeout_in_minutes: 45
depends_on:
- image-build-amd
source_file_dependencies:
+9 -13
View File
@@ -28,9 +28,8 @@ steps:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 40
device: mi325_1
timeout_in_minutes: 50
depends_on:
- image-build-amd
@@ -40,21 +39,19 @@ 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
timeout_in_minutes: 45
device: mi325_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 53
timeout_in_minutes: 35
device: h200_18gb
source_file_dependencies:
- vllm/v1/
@@ -63,8 +60,8 @@ steps:
- pytest -v -s v1/e2e/general/test_async_scheduling.py
mirror:
amd:
device: mi250_1
timeout_in_minutes: 55
device: mi325_1
timeout_in_minutes: 70
depends_on:
- image-build-amd
@@ -79,8 +76,8 @@ steps:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
mirror:
amd:
device: mi250_1
timeout_in_minutes: 50
device: mi325_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
@@ -119,7 +116,6 @@ steps:
amd:
dind: false
device: mi300_2
timeout_in_minutes: 30
depends_on:
- image-build-amd
+13 -20
View File
@@ -30,16 +30,16 @@ steps:
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 55
device: mi325_1
# TODO(akaratza): Test after Torch >= 2.12 bump
soft_fail: true
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server)
key: entrypoints-integration-api-server
device: h200_35gb
timeout_in_minutes: 75
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -52,16 +52,14 @@ steps:
- pytest -v -s entrypoints/scale_out
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 65
device: mi325_1
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server OpenAI - Part 1)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 68
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -72,8 +70,7 @@ steps:
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
mirror:
amd:
dind: false
device: mi300_1
device: mi325_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
@@ -81,7 +78,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 2)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 83
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -93,9 +90,8 @@ steps:
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 70
device: mi325_1
timeout_in_minutes: 80
depends_on:
- image-build-amd
@@ -117,8 +113,7 @@ steps:
- pytest -v -s entrypoints/anthropic
mirror:
amd:
dind: false
device: mi300_1
device: mi325_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
@@ -161,7 +156,7 @@ steps:
- label: Entrypoints Integration (Pooling)
device: h200_35gb
key: entrypoints-integration-pooling
timeout_in_minutes: 75
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -181,9 +176,7 @@ steps:
- pytest -s entrypoints/openai/correctness/
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 30
device: mi325_1
depends_on:
- image-build-amd
source_file_dependencies:
@@ -18,7 +18,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 30
depends_on:
- image-build-amd
source_file_dependencies:
@@ -52,5 +51,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
+4 -44
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
@@ -116,8 +80,7 @@ steps:
parallelism: 2
mirror:
amd:
dind: false
device: mi300_1
device: mi325_1
timeout_in_minutes: 90
depends_on:
- image-build-amd
@@ -155,9 +118,7 @@ steps:
parallelism: 2
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 120
device: mi325_1
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
@@ -187,9 +148,8 @@ steps:
parallelism: 5
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 55
device: mi325_1
timeout_in_minutes: 65
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
+5 -6
View File
@@ -14,9 +14,8 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 45
device: mi325_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
@@ -142,7 +141,7 @@ steps:
amd:
dind: false
device: mi300_8
timeout_in_minutes: 40
timeout_in_minutes: 60
depends_on:
- image-build-amd
commands:
@@ -337,7 +336,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 +346,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
+3 -5
View File
@@ -14,11 +14,9 @@ steps:
parallelism: 4
mirror:
amd:
dind: false
device: mi300_1
soft_fail: true
device: mi325_1
working_dir: "/vllm-workspace/tests"
timeout_in_minutes: 85
timeout_in_minutes: 65
source_file_dependencies:
- vllm/lora
- tests/lora
@@ -48,4 +46,4 @@ steps:
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
- pytest -v -s -x lora/test_olmoe_tp.py
- pytest -v -s -x lora/test_gptoss_tp.py
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+20 -36
View File
@@ -25,13 +25,13 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 50
timeout_in_minutes: 75
depends_on:
- image-build-amd
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 83
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/config/
@@ -59,9 +59,7 @@ steps:
- pytest -v -s v1/test_outputs.py
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 70
device: mi325_1
depends_on:
- image-build-amd
@@ -92,7 +90,6 @@ steps:
- tests/v1/kv_offload
- tests/v1/simple_kv_offload
- tests/v1/worker
- tests/v1/streaming_input
- tests/v1/kv_connector/unit
- tests/v1/ec_connector/unit
- tests/v1/metrics
@@ -106,7 +103,6 @@ steps:
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
- pytest -v -s v1/worker
- pytest -v -s v1/streaming_input
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
@@ -115,9 +111,8 @@ steps:
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 65
device: mi325_1
timeout_in_minutes: 75
depends_on:
- image-build-amd
@@ -148,8 +143,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
@@ -213,7 +206,7 @@ steps:
- vllm/multimodal
- examples/
commands:
- pip install --no-deps tensorizer # for tensorizer test
- pip install tensorizer # for tensorizer test
# for basic
- python3 basic/offline_inference/chat.py
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -237,9 +230,7 @@ steps:
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 75
device: mi325_1
source_file_dependencies:
- vllm/entrypoints
- vllm/multimodal
@@ -266,7 +257,6 @@ steps:
- vllm/utils/
- vllm/v1/
- tests/v1/tracing
- tests/tracing/
commands:
- "pip install \
'opentelemetry-sdk>=1.26.0' \
@@ -274,14 +264,12 @@ 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
device: mi300_2
timeout_in_minutes: 30
device: mi325_2
depends_on:
- image-build-amd
optional: true
- label: Python-only Installation
key: python-only-installation
@@ -296,9 +284,8 @@ steps:
- bash standalone_tests/python_only_compile.sh
mirror:
amd:
device: mi250_1
timeout_in_minutes: 55
soft_fail: true
device: mi325_1
timeout_in_minutes: 45
depends_on:
- image-build-amd
source_file_dependencies:
@@ -398,7 +385,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 +395,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 +410,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 +426,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 -2
View File
@@ -5,7 +5,7 @@ steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 60
timeout_in_minutes: 45
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
@@ -30,7 +30,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
+1 -1
View File
@@ -41,7 +41,7 @@ steps:
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pip install --no-deps tensorizer # for tensorizer test
- pip install tensorizer # for tensorizer test
- python3 basic/offline_inference/chat.py # for basic
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
+2 -30
View File
@@ -42,37 +42,10 @@ steps:
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 50
device: mi325_1
depends_on:
- image-build-amd
- label: Inkling Unit Tests (B200)
key: inkling-unit-tests-b200
timeout_in_minutes: 40
device: b200-k8s
source_file_dependencies:
- vllm/models/inkling/
- vllm/cute_utils/
- cmake/external_projects/tml_fa4.cmake
- tests/models/inkling/
commands:
# 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/
commands:
# The native NVIDIA AttnRes kernel requires the SM100 family.
- pytest -v -s models/kimi_k3
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
@@ -82,8 +55,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/
+7 -9
View File
@@ -17,7 +17,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 45
depends_on:
- image-build-amd
@@ -40,7 +39,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 40
depends_on:
- image-build-amd
source_file_dependencies:
@@ -63,6 +61,7 @@ steps:
- tests/models/language/generation
commands:
# Install fast path packages for testing against transformers
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
# Shard the hybrid language model tests that are numerically stable on Hopper.
@@ -70,9 +69,8 @@ steps:
parallelism: 2
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 60
device: mi325_1
timeout_in_minutes: 70
depends_on:
- image-build-amd
commands:
@@ -104,6 +102,7 @@ steps:
- tests/models/language/generation
commands:
# Install fast path packages for testing against transformers
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
@@ -131,15 +130,14 @@ steps:
- pytest -v -s models/language/pooling -m 'not core_model'
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 95
device: mi325_1
timeout_in_minutes: 120
depends_on:
- image-build-amd
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
optional: true
source_file_dependencies:
+11 -22
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -14,15 +14,13 @@ steps:
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 65
device: mi325_1
depends_on:
- image-build-amd
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 75
timeout_in_minutes: 50
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -33,9 +31,7 @@ steps:
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 55
device: mi325_1
depends_on:
- image-build-amd
@@ -51,15 +47,14 @@ steps:
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
mirror:
amd:
device: mi250_1
timeout_in_minutes: 55
device: mi325_1
depends_on:
- image-build-amd
- label: "Multi-Modal Models (Standard) 4: other + whisper"
device: h200_35gb
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 75
timeout_in_minutes: 50
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -70,9 +65,7 @@ steps:
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 50
device: mi325_1
depends_on:
- image-build-amd
@@ -92,7 +85,7 @@ steps:
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 98
timeout_in_minutes: 65
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -116,7 +109,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 35
depends_on:
- image-build-amd
source_file_dependencies:
@@ -139,9 +131,7 @@ steps:
- pytest -v -s models/multimodal/test_mapping.py
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 90
device: mi325_1
depends_on:
- image-build-amd
@@ -176,9 +166,8 @@ steps:
- pytest -v -s models/multimodal/pooling -m 'not core_model'
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 60
device: mi325_1
timeout_in_minutes: 75
depends_on:
- image-build-amd
source_file_dependencies:
+8 -2
View File
@@ -5,7 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
device: h200_35gb
key: pytorch-compilation-unit-tests
timeout_in_minutes: 150
timeout_in_minutes: 90
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
@@ -107,6 +107,13 @@ steps:
- tests/compile/passes
commands:
- pytest -s -v compile/passes --ignore compile/passes/distributed
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
- label: PyTorch Fullgraph Smoke Test
device: h200_35gb
@@ -229,7 +236,6 @@ steps:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 30
depends_on:
- image-build-amd
source_file_dependencies:
+4 -2
View File
@@ -24,7 +24,8 @@ steps:
- uv pip install --system conch-triton-kernels
# The SM90-only checkpoint currently contains a removed weight_chan_scale
# parameter. It was not exercised by the previous L4 job.
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 8
- label: Quantized Fusions
device: h200_35gb
@@ -67,4 +68,5 @@ steps:
- vllm/model_executor/layers/quantization
- tests/models/quantization
commands:
- pytest -v -s models/quantization
- pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 3
+30
View File
@@ -8,6 +8,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/benchmarks/
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
@@ -23,6 +28,7 @@ steps:
- tests/entrypoints/openai/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- 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)"
@@ -43,6 +49,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
- vllm/v1/engine/
@@ -54,6 +65,7 @@ steps:
# - tests/entrypoints/serve/dev/test_sleep.py
- tests/entrypoints/serve/tokenize/test_tokenization.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
@@ -72,10 +84,16 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- tests/utils.py
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
@@ -86,11 +104,17 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/tool_parsers/
- tests/utils.py
- tests/tool_use/
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
@@ -101,6 +125,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/distributed/
- vllm/engine/
- vllm/executor/
@@ -111,6 +140,7 @@ steps:
- tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
@@ -26,5 +26,7 @@ steps:
- rust-toolchain.toml
- .buildkite/test_areas/rust_frontend_cargo.yaml
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
commands:
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
+1 -11
View File
@@ -19,18 +19,8 @@ steps:
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
mirror:
amd:
device: mi250_1
timeout_in_minutes: 40
device: mi325_1
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/model_executor/layers
- vllm/sampling_metadata.py
- vllm/v1/sample/
- vllm/entrypoints/generate/beam_search/
- tests/samplers
- tests/conftest.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s samplers
+9 -41
View File
@@ -14,9 +14,8 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 55
device: mi325_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
@@ -54,9 +53,8 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 75
device: mi325_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
source_file_dependencies:
@@ -90,15 +88,14 @@ 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
device: mi300_1
timeout_in_minutes: 35
device: mi325_1
timeout_in_minutes: 55
# TODO(akaratza): Test after Torch >= 2.12 bump
soft_fail: true
depends_on:
- image-build-amd
source_file_dependencies:
@@ -122,8 +119,7 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
mirror:
amd:
dind: false
device: mi300_1
device: mi325_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
@@ -174,31 +170,3 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
key: spec-decode-deepseek-mtp-parallel-load-b200
timeout_in_minutes: 30
device: b200-k8s
optional: true
num_devices: 2
source_file_dependencies:
- vllm/v1/spec_decode/llm_base_proposer.py
- vllm/v1/spec_decode/eagle.py
- vllm/v1/worker/gpu/spec_decode/eagle/
- vllm/model_executor/models/deepseek_mtp.py
- vllm/model_executor/models/deepseek_v2.py
- 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
@@ -17,7 +17,6 @@ steps:
amd:
dind: false
device: mi300_2
timeout_in_minutes: 35
depends_on:
- image-build-amd
commands:
-38
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:
@@ -182,18 +181,6 @@ pull_request_rules:
add:
- performance
- name: label-quantization
description: Automatically apply quantization label
conditions:
- label != stale
- or:
- files~=^vllm/model_executor/layers/quantization/
- title~=(?i)quant
actions:
label:
add:
- quantization
- name: label-qwen
description: Automatically apply qwen label
conditions:
@@ -233,31 +220,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:
+1 -61
View File
@@ -130,66 +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: [
{
term: "quantization",
searchIn: "both"
},
{
term: "quantized",
searchIn: "both"
},
],
},
"intel-gpu": {
// Keyword search - matches whole words only (with word boundaries)
keywords: [
{
term: "B50",
searchIn: "both"
},
{
term: "B60",
searchIn: "both"
},
{
term: "B70",
searchIn: "both"
},
{
term: "intel gpu",
searchIn: "both"
},
{
term: "Arc GPU",
searchIn: "both"
},
{
term: "BMG",
searchIn: "both"
},
],
},
// Add more label configurations here as needed
// example: {
// keywords: [...],
@@ -551,4 +491,4 @@ jobs:
issue_number: context.issue.number,
body: message,
});
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
+1 -1
View File
@@ -29,7 +29,7 @@ jobs:
steps:
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
- uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
- uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
enable-cache: true
cache-dependency-glob: |
+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: read
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@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.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 }}
-580
View File
@@ -1,580 +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"]],
"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,362 +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"],
"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()
+4
View File
@@ -173,6 +173,9 @@ venv.bak/
# mkdocs documentation
/site
docs/argparse
docs/examples/*
!docs/examples/README.md
# mypy
.mypy_cache/
@@ -254,3 +257,4 @@ vllm/grpc/vllm_engine_pb2.pyi
# Ignore generated cpu headers
csrc/cpu/cpu_attn_dispatch_generated.h
rust-coverage-tools/
-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
+13 -45
View File
@@ -68,8 +68,8 @@ endif()
# requirements.txt files and should be kept consistent. The ROCm torch
# versions are derived from docker/Dockerfile.rocm
#
set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
# version check would always warn. Only treat it as a nightly build when the
# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
@@ -114,11 +114,6 @@ find_package(Torch REQUIRED)
# Supported NVIDIA architectures.
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
# Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
# targeting it directly.
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
# to support the whole generation without specifying all sub-architectures
@@ -219,8 +214,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 +227,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
@@ -430,7 +420,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
"9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
@@ -705,7 +695,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
@@ -825,7 +815,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
# require CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -909,7 +899,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -934,7 +924,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# moe_data.cu is used by all CUTLASS MoE kernels.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
@@ -991,7 +981,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
# only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -1057,7 +1047,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# Runtime dispatch is gated in
# vllm/v1/attention/backends/mla/cutlass_mla.py.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -1079,24 +1069,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set(MLA_ARCHS)
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
"10.0f" "${CUDA_ARCHS}")
endif()
if(KIMI_K3_ATTN_RES_ARCHS)
set(KIMI_K3_ATTN_RES_SRC
"csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${KIMI_K3_ATTN_RES_SRC}"
CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
COMPILE_OPTIONS
"$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
message(STATUS
"Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
endif()
# Hadacore kernels
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
if(HADACORE_ARCHS)
@@ -1138,10 +1110,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
if(KIMI_K3_ATTN_RES_ARCHS)
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
endif()
# Needed by CUTLASS kernels
target_compile_definitions(_C_stable_libtorch PRIVATE
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -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())
+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,
-267
View File
@@ -1,267 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
long greedy decode (CUDA graphs on) once with the standard kernel and once with
ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
separate subprocesses so each gets a clean CUDA context.
The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
CUDA-graph capture on the pre-release Blackwell FP4 path); pass
--no-disable-flashinfer-autotune for non-FP4 models.
Examples:
python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
--model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4
"""
import argparse
import json
import os
import subprocess
import sys
import time
DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
def parse_args():
p = argparse.ArgumentParser(
description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
)
p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
p.add_argument("--prompt", default=DEFAULT_PROMPT)
p.add_argument("--batch-size", type=int, default=256)
p.add_argument("--num-steps", type=int, default=1000)
p.add_argument("--warmup-steps", type=int, default=128)
p.add_argument("--repeats", type=int, default=1)
p.add_argument(
"--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
)
p.add_argument(
"--dtype",
default="bfloat16",
choices=["bfloat16", "float16", "float32", "auto"],
)
p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
p.add_argument("--max-model-len", type=int, default=None)
p.add_argument(
"--disable-flashinfer-autotune",
action=argparse.BooleanOptionalAction,
default=True,
help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
"It is unstable under CUDA-graph capture on the "
"pre-release Blackwell FP4 path; pass "
"--no-disable-flashinfer-autotune for non-FP4 models.",
)
p.add_argument(
"--mamba-ssm-cache-dtype",
default="auto",
choices=["auto", "float32", "float16", "bfloat16"],
help="SSM state dtype (both modes). 'auto' = config-driven; "
"'float32' = fp32 state, 'bfloat16' = s16 state.",
)
p.add_argument(
"--baseline-ssm-config",
default="",
help="Pin the STANDARD baseline's SSM launch config as "
"'bsm,nw' via override_ssm_config (forces the in-process "
"engine so the override reaches the kernel). Empty = off.",
)
p.add_argument(
"--worker",
choices=["standard", "replayssm"],
default=None,
help=argparse.SUPPRESS,
)
return p.parse_args()
def resolve_max_model_len(args) -> int:
if args.max_model_len is not None:
return args.max_model_len
return args.num_steps + 256
def run_worker(args):
# override_ssm_config is a module global; it only reaches the model if the
# engine runs in-process (default V1 spawns a separate EngineCore). Force it.
if args.worker == "standard" and args.baseline_ssm_config:
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
import torch
from vllm import LLM, SamplingParams
mode = args.worker
max_model_len = resolve_max_model_len(args)
llm_kwargs = dict(
model=args.model_id,
tensor_parallel_size=1,
dtype=args.dtype,
max_model_len=max_model_len,
trust_remote_code=True,
enable_prefix_caching=False,
enable_chunked_prefill=False,
max_num_seqs=args.batch_size,
max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
enforce_eager=False,
disable_log_stats=True,
gpu_memory_utilization=args.gpu_memory_utilization,
# SSM state dtype (applies to both standard and ReplaySSM).
mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
)
if args.disable_flashinfer_autotune:
# FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
# re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
if mode == "replayssm":
llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
_ssm_cm = None
if mode == "standard" and args.baseline_ssm_config:
from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
_bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
_ssm_cm = override_ssm_config((_bsm, _nw))
_ssm_cm.__enter__() # active through LLM() graph capture + decode
print(
f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
flush=True,
)
llm = LLM(**llm_kwargs)
prompts = [args.prompt] * args.batch_size
def timed_generate(n_tokens):
sp = SamplingParams(
n=1,
temperature=0.0,
ignore_eos=True,
min_tokens=n_tokens,
max_tokens=n_tokens,
)
if torch.accelerator.is_available():
torch.accelerator.synchronize()
t0 = time.perf_counter()
outs = llm.generate(prompts, sp, use_tqdm=False)
if torch.accelerator.is_available():
torch.accelerator.synchronize()
elapsed = time.perf_counter() - t0
produced = min(len(o.outputs[0].token_ids) for o in outs)
assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
return elapsed
timed_generate(args.warmup_steps)
best = None
for _ in range(args.repeats):
elapsed = timed_generate(args.num_steps)
tok_s = args.batch_size * args.num_steps / elapsed
per_step_ms = elapsed / args.num_steps * 1e3
print(
f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step",
flush=True,
)
if best is None or elapsed < best["elapsed_s"]:
best = {
"mode": mode,
"elapsed_s": elapsed,
"tok_s": tok_s,
"per_step_ms": per_step_ms,
}
print("RESULT_JSON " + json.dumps(best), flush=True)
if _ssm_cm is not None:
_ssm_cm.__exit__(None, None, None)
def run_one_mode(args, mode) -> dict:
cmd = [
sys.executable,
__file__,
"--worker",
mode,
"--model-id",
args.model_id,
"--prompt",
args.prompt,
"--batch-size",
str(args.batch_size),
"--num-steps",
str(args.num_steps),
"--warmup-steps",
str(args.warmup_steps),
"--repeats",
str(args.repeats),
"--buffer-len",
str(args.buffer_len),
"--dtype",
args.dtype,
"--gpu-memory-utilization",
str(args.gpu_memory_utilization),
"--mamba-ssm-cache-dtype",
args.mamba_ssm_cache_dtype,
"--baseline-ssm-config",
args.baseline_ssm_config,
]
cmd.append(
"--disable-flashinfer-autotune"
if args.disable_flashinfer_autotune
else "--no-disable-flashinfer-autotune"
)
if args.max_model_len is not None:
cmd += ["--max-model-len", str(args.max_model_len)]
result = None
proc = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
)
for line in proc.stdout:
sys.stdout.write(line)
sys.stdout.flush()
if line.startswith("RESULT_JSON "):
result = json.loads(line[len("RESULT_JSON ") :])
proc.wait()
if proc.returncode != 0:
raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
if result is None:
raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
return result
def main():
args = parse_args()
if args.worker is not None:
run_worker(args)
return
print(
f"model={args.model_id} batch_size={args.batch_size} "
f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}"
)
std = run_one_mode(args, "standard")
fla = run_one_mode(args, "replayssm")
speedup = std["per_step_ms"] / fla["per_step_ms"]
print()
header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
print(header)
print("-" * len(header))
for r in (std, fla):
print(
f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
)
print("-" * len(header))
print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
if __name__ == "__main__":
main()
+37
View File
@@ -8,6 +8,8 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
CARGO_LLVM_COV_VERSION="0.8.7"
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -30,4 +32,39 @@ else
PROFILE_ARG="--release"
fi
rm -rf "$COVERAGE_TOOLS_DIR"
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
# Keep instrumented intermediates isolated when local builds switch modes.
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
cargo +"$TOOLCHAIN" install \
--locked \
--version "$CARGO_LLVM_COV_VERSION" \
cargo-llvm-cov
eval "$(
cargo +"$TOOLCHAIN" llvm-cov show-env \
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
--sh
)"
# Build scripts and proc macros can run during compilation. Their profiles
# are unrelated to runtime coverage and would otherwise pollute the tree.
export LLVM_PROFILE_FILE=/dev/null
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
fi
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
--print target-libdir)")/bin"
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
fi
+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)
-3
View File
@@ -68,9 +68,6 @@ endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
endif()
else()
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
endif()
+1 -3
View File
@@ -60,9 +60,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
# CUDA 12.9 has introduced "Family-Specific Architecture Features"
# this supports all compute_10x family
list(APPEND SUPPORT_ARCHS "10.0f")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
list(APPEND SUPPORT_ARCHS "10.7f")
endif()
elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
list(APPEND SUPPORT_ARCHS "10.0a")
endif()
@@ -191,3 +188,4 @@ else()
add_custom_target(_flashmla_C)
add_custom_target(_flashmla_extension_C)
endif()
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
+1 -5
View File
@@ -55,11 +55,7 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
endif()
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+10 -21
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()
@@ -397,24 +396,14 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
# match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
# family. The output uses TGT's value to preserve the user's compilation flags.
set(_CUDA_ARCHS)
# Resolve exact base matches before family fallbacks so a generic entry such
# as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
foreach(_arch ${_SRC_CUDA_ARCHS})
if(_arch MATCHES "[af]$")
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
list(APPEND _CUDA_ARCHS "${_arch}")
endif()
endif()
endforeach()
foreach(_arch ${_SRC_CUDA_ARCHS})
if(_arch MATCHES "[af]$")
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
list(APPEND _CUDA_ARCHS "${_arch}")
elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
list(APPEND _CUDA_ARCHS "${_base}a")
elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
@@ -498,7 +487,7 @@ endfunction()
function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
endif()
+13
View File
@@ -10,3 +10,16 @@ fixes:
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
- "/usr/lib/python3.*/site-packages/vllm/::vllm/"
# Map Rust sources built in the E2E image and on Buildkite agents.
- "/workspace/rust/::rust/"
- "/var/lib/buildkite-agent/.*/rust/::rust/"
flags:
rust-unit:
paths:
- rust/
carryforward: false
rust-e2e:
paths:
- rust/
carryforward: false
-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
+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]);
@@ -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});
+5 -34
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);
@@ -207,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,
@@ -516,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
@@ -611,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) "
"-> ()");
@@ -622,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,"
+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},
};
-3
View File
@@ -1025,9 +1025,6 @@ __global__ void gather_and_maybe_dequant_cache(
batch_offset += offset;
int32_t block_table_id = batch_offset / block_size;
int32_t slot_id = batch_offset % block_size;
// seq_starts may push the block index past the end of the batch's block
// table row.
if (block_table_id >= block_table_stride) continue;
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
int32_t block_id = block_table[block_table_offset];
int64_t cache_offset =
@@ -1,954 +0,0 @@
/*
* Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
*/
// Production AttnRes forward for Blackwell (SM100).
//
// Warp-specialized online softmax + residual + RMSNorm:
// - 1 producer warp issues cp.async.bulk row loads into shared memory.
// - 8 consumer warps compute reductions and output.
// - Q=res_weight*rms_weight remains in registers across persistent tokens.
// - V rows are converted once and cached as FP32 in TMEM between passes.
//
// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major
// block residual storage.
#include "../torch_utils.h"
#include <cfloat>
#include <cstdint>
#include <cstdio>
#include <cuda_runtime.h>
#include <type_traits>
using bf16_t = __nv_bfloat16;
namespace sm100 {
namespace fwd_prod_v2 {
constexpr int K_TILE = 1024;
constexpr int N_CHUNK_DEFAULT = 4;
constexpr int CHUNK_DEPTH = 2;
constexpr int BLK = 288; // 1 producer warp + 8 consumer warps
constexpr int CONSUMER_THREADS = BLK - 32; // 256
constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32;
constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups
constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS;
constexpr int FIRST_USER_NAMED_BARRIER = 8;
__device__ __forceinline__ const bf16_t* residual_addr(
const bf16_t* block_res, const bf16_t* layer_res, int source, int N,
int token, int block_stride_m, int block_stride_r, int H) {
if (source < N - 1) {
return block_res + static_cast<long long>(token) * block_stride_m +
source * block_stride_r;
}
return layer_res + static_cast<long long>(token) * H;
}
__device__ __forceinline__ uint32_t elect_one_sync() {
uint32_t pred = 0;
uint32_t laneid = 0;
asm volatile(
"{\n"
".reg .b32 %%rx;\n"
".reg .pred %%px;\n"
" elect.sync %%rx|%%px, %2;\n"
"@%%px mov.s32 %1, 1;\n"
" mov.s32 %0, %%rx;\n"
"}\n"
: "+r"(laneid), "+r"(pred)
: "r"(0xffffffff));
return pred;
}
__device__ __forceinline__ void mbarrier_init(uint64_t& barrier,
int thread_count) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr),
"r"(thread_count));
}
__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier,
uint32_t bytes) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(
barrier_addr),
"r"(bytes));
}
__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile(
"{\n"
".reg .pred p;\n"
"WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n"
"@p bra DONE;\n"
"bra WAIT;\n"
"DONE:\n"
"}\n" ::"r"(barrier_addr),
"r"(phase));
}
__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile(
"{\n"
".reg .b64 state;\n"
"mbarrier.arrive.shared::cta.b64 state, [%0];\n"
"}\n" ::"r"(barrier_addr));
}
__device__ __forceinline__ void fence_mbarrier_init() {
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads,
uint32_t user_barrier_id) {
asm volatile(
"bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER),
"r"(num_threads)
: "memory");
}
__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) {
uint32_t const dst_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(dst));
asm volatile(
"tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
dst_addr),
"r"(num_columns));
}
__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) {
asm volatile(
"tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr),
"r"(num_columns));
}
__device__ __forceinline__ void tmem_release_allocation_lock() {
asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
}
__device__ __forceinline__ void tmem_store_wait() {
asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
}
template <int N, typename T>
__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) {
uint32_t* values = reinterpret_cast<uint32_t*>(dst);
if constexpr (N == 8) {
asm volatile(
"tcgen05.ld.sync.aligned.32x32b.x8.b32"
"{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
: "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]),
"=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7])
: "r"(src_addr));
} else {
static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
asm volatile(
"tcgen05.ld.sync.aligned.32x32b.x4.b32"
"{%0, %1, %2, %3}, [%4];\n"
: "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3])
: "r"(src_addr));
}
}
template <int N, typename T>
__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) {
uint32_t* values = reinterpret_cast<uint32_t*>(src);
if constexpr (N == 8) {
asm volatile(
"tcgen05.st.sync.aligned.32x32b.x8.b32"
"[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]),
"r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]),
"r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr));
} else {
static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
asm volatile(
"tcgen05.st.sync.aligned.32x32b.x4.b32"
"[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]),
"r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr));
}
}
__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) {
float2 result;
asm volatile("add.rn.f32x2 %0, %1, %2;\n"
: "=l"(reinterpret_cast<uint64_t&>(result))
: "l"(reinterpret_cast<uint64_t const&>(a)),
"l"(reinterpret_cast<uint64_t const&>(b)));
return result;
}
__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) {
float2 result;
asm volatile("mul.f32x2 %0, %1, %2;\n"
: "=l"(reinterpret_cast<uint64_t&>(result))
: "l"(reinterpret_cast<uint64_t const&>(a)),
"l"(reinterpret_cast<uint64_t const&>(b)));
return result;
}
__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b,
const float2& c) {
float2 result;
asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
: "=l"(reinterpret_cast<uint64_t&>(result))
: "l"(reinterpret_cast<uint64_t const&>(a)),
"l"(reinterpret_cast<uint64_t const&>(b)),
"l"(reinterpret_cast<uint64_t const&>(c)));
return result;
}
template <int NC>
struct FwdSmemPlan {
alignas(16) uint64_t bar_ready[CHUNK_DEPTH];
alignas(16) uint64_t bar_consumed[CHUNK_DEPTH];
alignas(16) uint64_t bar_output_norm_ready;
alignas(16) float2 ws_stats[CONSUMER_WARPS][NC];
uint32_t tmem_base;
};
__device__ __forceinline__ void cp_async_bulk(void* smem_dst,
const void* gmem_src, int bytes,
uint64_t& mbar) {
uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst));
uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar));
asm volatile(
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], "
"[%1], %2, [%3];\n" ::"r"(s),
"l"(gmem_src), "r"(bytes), "r"(m)
: "memory");
}
template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
bool OUTPUT_NORM_IN_SMEM = false>
__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T,
int B, int block_stride_m, int block_stride_r, float rms_eps,
const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
constexpr float LOG2_E = 1.4426950408889634f;
constexpr int N_CHUNK = NC;
// The two-source specialization only consumes half of the TMEM columns.
constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256;
constexpr int NUM_BUFS = CHUNK_DEPTH * NC;
constexpr int NHT = H / K_TILE;
constexpr int SLICES_PER_GROUP =
(NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS;
constexpr int VEC = 8;
constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC;
constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC;
constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP;
static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC);
static_assert(H >= 4096 && H <= 8192);
static_assert(H % K_TILE == 0);
const int tid = threadIdx.x;
const int wid = tid >> 5;
const int lane = tid & 31;
const int TB = T * B;
const int num_ctas = gridDim.x;
const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
const int comp_wid = wid - 1;
const int comp_tid = tid - 32;
const int group = (comp_wid >= 4) ? 1 : 0;
const int ct_in_group =
(comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1;
const int k_local = ct_in_group * VEC;
constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t);
constexpr size_t DELTA_BYTES =
HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0;
constexpr size_t OUTPUT_NORM_BYTES =
OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0;
extern __shared__ __align__(16) char smem_raw[];
bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw); // [NUM_BUFS][H]
bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES);
bf16_t* output_norm_buf =
reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES);
FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>(
smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES);
auto slot_of = [](long long gci, int n) {
return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n;
};
auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); };
auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; };
auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* {
return delta_bufs + chunk_slot * H;
};
if (wid == 0 && elect_one_sync()) {
#pragma unroll
for (int i = 0; i < CHUNK_DEPTH; i++) {
mbarrier_init(plan.bar_ready[i], 1);
mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS);
}
if constexpr (OUTPUT_NORM_IN_SMEM) {
mbarrier_init(plan.bar_output_norm_ready, 1);
}
fence_mbarrier_init();
}
// gdc wait BEFORE tmem alloc
cudaGridDependencySynchronize();
if (wid == 1) {
tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base);
if constexpr (RELEASE_TMEM) {
tmem_release_allocation_lock();
}
}
__syncthreads();
if constexpr (OUTPUT_NORM_IN_SMEM) {
if (wid == 0 && elect_one_sync()) {
mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t));
cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t),
plan.bar_output_norm_ready);
}
}
const uint32_t my_v_tmem =
comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0;
float q_cache[ACC_PER_THREAD];
if (comp_tid >= 0) {
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
#pragma unroll
for (int j = 0; j < 4; j++) {
int h = h_base + j;
q_cache[si * VEC + j] =
__bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
#pragma unroll
for (int j = 0; j < VEC; j++) {
int h = h_base + j;
q_cache[si * VEC + j] =
__bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
}
}
}
if (wid == 0) {
if (elect_one_sync()) {
long long gci = 0;
for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
const int t = tb / B;
for (int ci = 0; ci < num_chunks; ci++, gci++) {
int ns = ci * N_CHUNK;
int an = min(N_CHUNK, N - ns);
int chunk_slot = (int)(gci % CHUNK_DEPTH);
int pc = phase_of(gci);
mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1);
int transaction_bytes = an * H * (int)sizeof(bf16_t);
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (prefix_n >= 0 && prefix_n < an) {
transaction_bytes += H * (int)sizeof(bf16_t);
}
}
mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes);
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
if (n >= an) continue;
int slot = slot_of(gci, n);
const bf16_t* src =
residual_addr(block_res, layer_res, ns + n, N, t,
block_stride_m, block_stride_r, H);
cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t),
plan.bar_ready[chunk_slot]);
}
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (prefix_n >= 0 && prefix_n < an) {
cp_async_bulk(delta_buf_ptr(chunk_slot),
delta + (long long)tb * H, H * sizeof(bf16_t),
plan.bar_ready[chunk_slot]);
}
}
}
}
}
} else {
float acc32[ACC_PER_THREAD] = {};
float eps_cache;
asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps));
long long gci = 0;
for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
float m_running = -FLT_MAX;
float s_running = 0.f;
#pragma unroll
for (int i = 0; i < ACC_PER_THREAD; i++) {
acc32[i] = 0.f;
}
for (int ci = 0; ci < num_chunks; ci++, gci++) {
int ns = ci * N_CHUNK;
int an = min(N_CHUNK, N - ns);
int chunk_slot = (int)(gci % CHUNK_DEPTH);
int pr = phase_of(gci);
mbarrier_wait(plan.bar_ready[chunk_slot], pr);
float2 sq_local[N_CHUNK] = {};
float2 dot_local[N_CHUNK] = {};
auto pass_A_body = [&](auto AN_TOK) {
constexpr int AN = decltype(AN_TOK)::value;
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base =
6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
const float* qv = &q_cache[si * VEC];
#pragma unroll
for (int n = 0; n < AN; n++) {
int slot = slot_of(gci, n);
int2 vp =
*reinterpret_cast<const int2*>(buf_ptr(slot) + h_base);
auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (n == prefix_n) {
const bf16_t* delta_ptr =
delta_buf_ptr(chunk_slot) + h_base;
#pragma unroll
for (int j = 0; j < 2; j++) {
auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
delta_ptr + 2 * j);
v2[j] = __hadd2(v2[j], delta2);
}
*reinterpret_cast<int2*>(layer_res + (long long)tb * H +
h_base) = vp;
}
}
float2 f[2] = {__bfloat1622float2(v2[0]),
__bfloat1622float2(v2[1])};
tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
sq_local[n] = float2_fma(f[0], f[0], sq_local[n]);
sq_local[n] = float2_fma(f[1], f[1], sq_local[n]);
dot_local[n] =
float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]);
dot_local[n] =
float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]);
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
const float* qv = &q_cache[si * VEC];
#pragma unroll
for (int n = 0; n < AN; n++) {
int slot = slot_of(gci, n);
int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base);
auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (n == prefix_n) {
const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base;
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
delta_ptr + 2 * j);
v2[j] = __hadd2(v2[j], delta2);
}
*reinterpret_cast<int4*>(layer_res + (long long)tb * H +
h_base) = vp;
}
}
float2 f[4] = {
__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]),
__bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])};
tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
sq_local[n] = float2_fma(f[j], f[j], sq_local[n]);
dot_local[n] = float2_fma(
f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]);
}
}
}
};
if constexpr (NC == 4) {
switch (an) {
case 4:
pass_A_body(std::integral_constant<int, 4>{});
break;
case 3:
pass_A_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_A_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_A_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else if constexpr (NC == 3) {
switch (an) {
case 3:
pass_A_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_A_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_A_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else {
static_assert(NC == 2);
switch (an) {
case 2:
pass_A_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_A_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
}
if (lane == 0) {
mbarrier_arrive(plan.bar_consumed[chunk_slot]);
}
tmem_store_wait();
float2 reduce_pair[N_CHUNK];
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y,
dot_local[n].x + dot_local[n].y);
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]);
packed = __shfl_xor_sync(0xffffffff, packed, offset);
float2 other = reinterpret_cast<float2&>(packed);
reduce_pair[n] = float2_add(reduce_pair[n], other);
}
}
if (lane == 0) {
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
plan.ws_stats[comp_wid][n] = reduce_pair[n];
}
}
named_barrier_sync(CONSUMER_THREADS, 0);
float local_rsig = 0.f;
float local_logit = 0.f;
int stat_n = lane / CONSUMER_WARPS;
int stat_w = lane % CONSUMER_WARPS;
float2 totals = {};
if (stat_n < N_CHUNK) {
totals = plan.ws_stats[stat_w][stat_n];
}
#pragma unroll
for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
totals.x +=
__shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS);
totals.y +=
__shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS);
}
if (stat_n < N_CHUNK && stat_w == 0) {
local_rsig = rsqrtf(totals.x / H + eps_cache);
local_logit = totals.y * local_rsig;
}
float logit_n[N_CHUNK];
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS);
}
float m_chunk = -FLT_MAX;
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]);
}
float m_new = fmaxf(m_running, m_chunk);
float corr = exp2f((m_running - m_new) * LOG2_E);
float w_n[N_CHUNK] = {};
float w_sum = 0.f;
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
if (n < an) {
w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E);
w_sum += w_n[n];
}
}
auto pass_B_body = [&](auto AN_TOK) {
constexpr int AN = decltype(AN_TOK)::value;
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
float2 corr2 = make_float2(corr, corr);
float2 a[2];
#pragma unroll
for (int j = 0; j < 2; j++) {
float2 old = make_float2(acc32[si * VEC + 2 * j],
acc32[si * VEC + 2 * j + 1]);
a[j] = float2_mul(old, corr2);
}
float2 f_cache[AN][2];
#pragma unroll
for (int n = 0; n < AN; n++) {
tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC,
f_cache[n]);
}
#pragma unroll
for (int n = 0; n < AN; n++) {
float2 wn = make_float2(w_n[n], w_n[n]);
#pragma unroll
for (int j = 0; j < 2; j++) {
a[j] = float2_fma(wn, f_cache[n][j], a[j]);
}
}
#pragma unroll
for (int j = 0; j < 2; j++) {
acc32[si * VEC + 2 * j] = a[j].x;
acc32[si * VEC + 2 * j + 1] = a[j].y;
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
float2 corr2 = make_float2(corr, corr);
float2 a[VEC / 2];
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
float2 old = make_float2(acc32[si * VEC + 2 * j],
acc32[si * VEC + 2 * j + 1]);
a[j] = float2_mul(old, corr2);
}
float2 f_cache[AN][VEC / 2];
#pragma unroll
for (int n = 0; n < AN; n++) {
tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]);
}
#pragma unroll
for (int n = 0; n < AN; n++) {
float2 wn = make_float2(w_n[n], w_n[n]);
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
a[j] = float2_fma(wn, f_cache[n][j], a[j]);
}
}
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
acc32[si * VEC + 2 * j] = a[j].x;
acc32[si * VEC + 2 * j + 1] = a[j].y;
}
}
};
if constexpr (NC == 4) {
switch (an) {
case 4:
pass_B_body(std::integral_constant<int, 4>{});
break;
case 3:
pass_B_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_B_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_B_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else if constexpr (NC == 3) {
switch (an) {
case 3:
pass_B_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_B_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_B_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else {
static_assert(NC == 2);
switch (an) {
case 2:
pass_B_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_B_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
}
s_running = s_running * corr + w_sum;
m_running = m_new;
}
float inv_s = 1.f / s_running;
bf16_t* out_ptr = output + (long long)tb * H;
float2 output_sq_pair = {};
// When output RMSNorm is fused, the softmax denominator cancels:
// (acc / s) * rsqrt(mean((acc / s)^2) + eps)
// = acc * rsqrt(mean(acc^2) + eps * s^2).
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
uint2 packed;
auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
float2 inv2 = make_float2(inv_s, inv_s);
#pragma unroll
for (int j = 0; j < 2; j++) {
float2 old = make_float2(acc32[si * VEC + 2 * j],
acc32[si * VEC + 2 * j + 1]);
if constexpr (HAS_OUTPUT_NORM) {
output_sq_pair = float2_fma(old, old, output_sq_pair);
} else {
float2 mixed = float2_mul(old, inv2);
ov2[j] = __float22bfloat162_rn(mixed);
}
}
if constexpr (!HAS_OUTPUT_NORM) {
*reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
uint4 packed;
auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
float2 inv2 = make_float2(inv_s, inv_s);
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
float2 old =
make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]);
if constexpr (HAS_OUTPUT_NORM) {
output_sq_pair = float2_fma(old, old, output_sq_pair);
} else {
float2 mixed = float2_mul(old, inv2);
ov2[j] = __float22bfloat162_rn(mixed);
}
}
if constexpr (!HAS_OUTPUT_NORM) {
*reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
}
}
if constexpr (HAS_OUTPUT_NORM) {
if constexpr (OUTPUT_NORM_IN_SMEM) {
// The immutable weight copy is acquired once, at its first use.
if (tb == blockIdx.x) {
mbarrier_wait(plan.bar_output_norm_ready, 0);
}
}
float output_sq = output_sq_pair.x + output_sq_pair.y;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset);
}
if (lane == 0) {
plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f);
}
named_barrier_sync(CONSUMER_THREADS, 0);
float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f;
#pragma unroll
for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
total_sq +=
__shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS);
}
if (lane == 0) {
total_sq =
rsqrtf(total_sq / H + output_norm_eps * s_running * s_running);
}
float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0);
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
uint2 packed;
auto* values = reinterpret_cast<bf16_t*>(&packed);
#pragma unroll
for (int j = 0; j < 4; j++) {
const bf16_t* weight_ptr =
OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
float weight = __bfloat162float(weight_ptr[h_base + j]);
values[j] = __float2bfloat16(acc32[si * VEC + j] *
output_rsigma * weight);
}
*reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
uint4 packed;
auto* values = reinterpret_cast<bf16_t*>(&packed);
#pragma unroll
for (int j = 0; j < VEC; j++) {
const bf16_t* weight_ptr =
OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
float weight = __bfloat162float(weight_ptr[h_base + j]);
values[j] =
__float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight);
}
*reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
}
}
}
}
cudaTriggerProgrammaticLaunchCompletion();
__syncthreads();
if (wid == 1) {
tmem_free(plan.tmem_base, TMEM_COLS_ALLOC);
}
#else
if (threadIdx.x == 0) {
printf("attn_res_fwd_online_v2_kernel requires sm_10x\n");
}
#endif
}
template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
bool OUTPUT_NORM_IN_SMEM = false>
static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
const bf16_t* delta, const bf16_t* res_weight,
const bf16_t* rms_weight, bf16_t* output, int N, int T,
int B, float rms_eps, int num_sm, cudaStream_t stream,
const bf16_t* output_norm_weight = nullptr,
float output_norm_eps = 0.f, int block_stride_m = 0,
int block_stride_r = 0) {
constexpr size_t smem_size =
((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) +
(OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) +
sizeof(FwdSmemPlan<NC>) + 15) &
~size_t(15);
auto kernel =
&attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA,
HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
static bool attrs_set = false;
if (!attrs_set) {
if (smem_size > 48 * 1024) {
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
}
attrs_set = true;
}
int grid = RELEASE_TMEM ? num_sm * 2 : num_sm;
cudaLaunchConfig_t config{};
config.gridDim = grid;
config.blockDim = BLK;
config.dynamicSmemBytes = smem_size;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.attrs = attrs;
config.numAttrs = 1;
cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
res_weight, rms_weight, output, N, T, B, block_stride_m,
block_stride_r, rms_eps, output_norm_weight,
output_norm_eps);
}
} // namespace fwd_prod_v2
} // namespace sm100
void kimi_k3_attn_res(torch::stable::Tensor& prefix,
torch::stable::Tensor const& delta,
torch::stable::Tensor const& blocks,
torch::stable::Tensor const& norm_weight,
torch::stable::Tensor const& qk_weight,
torch::stable::Tensor const& output_norm_weight,
torch::stable::Tensor& output, int64_t num_blocks,
double eps, double output_norm_eps) {
int const num_tokens = static_cast<int>(prefix.size(0));
int const device = prefix.get_device_index();
torch::stable::accelerator::DeviceGuard const device_guard(device);
cudaDeviceProp const* properties = get_device_prop();
STD_TORCH_CHECK(properties->major == 10,
"Kimi K3 AttnRes requires the SM100 family");
using namespace sm100::fwd_prod_v2;
// Two-source chunks and two resident CTAs are beneficial once setup is
// amortized by the long, full eight-block prefill workload.
if (num_blocks == 8 && num_tokens >= 4096) {
launch_fwd<7168, 2, true, true, true, true>(
static_cast<bf16_t const*>(blocks.data_ptr()),
static_cast<bf16_t*>(prefix.data_ptr()),
static_cast<bf16_t const*>(delta.data_ptr()),
static_cast<bf16_t const*>(qk_weight.data_ptr()),
static_cast<bf16_t const*>(norm_weight.data_ptr()),
static_cast<bf16_t*>(output.data_ptr()),
static_cast<int>(num_blocks) + 1, num_tokens, 1,
static_cast<float>(eps), properties->multiProcessorCount,
get_current_cuda_stream(device),
static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
static_cast<int>(blocks.stride(1)));
} else {
launch_fwd<7168, 4, false, true, true, true>(
static_cast<bf16_t const*>(blocks.data_ptr()),
static_cast<bf16_t*>(prefix.data_ptr()),
static_cast<bf16_t const*>(delta.data_ptr()),
static_cast<bf16_t const*>(qk_weight.data_ptr()),
static_cast<bf16_t const*>(norm_weight.data_ptr()),
static_cast<bf16_t*>(output.data_ptr()),
static_cast<int>(num_blocks) + 1, num_tokens, 1,
static_cast<float>(eps), properties->multiProcessorCount,
get_current_cuda_stream(device),
static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
static_cast<int>(blocks.stride(1)));
}
cudaError_t const error = cudaGetLastError();
STD_TORCH_CHECK(
error == cudaSuccess,
"Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error));
}
+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);
+3 -6
View File
@@ -9,16 +9,14 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding);
std::optional<torch::stable::Tensor> bias);
void topk_sigmoid(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding);
double routed_scaling_factor);
void topk_softplus_sqrt(
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
@@ -27,8 +25,7 @@ void topk_softplus_sqrt(
double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding);
const std::optional<torch::stable::Tensor>& tid2eid);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
@@ -174,8 +174,7 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int end_expert,
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
const bool* is_padding)
const double routed_scaling_factor)
{
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -229,14 +228,12 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int expert = result_kvp.key;
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
const int idx = k * block_row + k_idx;
// Return the unbiased scores for output weights
output[idx] = inputs_after_softmax[thread_read_offset + expert];
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : num_experts);
assert(is_pad_row || indices[idx] >= 0);
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
source_rows[idx] = k_idx * num_rows + block_row;
if (renormalize) {
selected_sum += inputs_after_softmax[thread_read_offset + expert];
@@ -280,7 +277,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor, const bool* is_padding)
const float* bias, const double routed_scaling_factor)
{
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -548,14 +545,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
// single) thread per row of the input/output matrices.
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -610,7 +605,7 @@ struct TopkConstants
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
{
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -621,7 +616,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
}
#ifndef USE_ROCM
@@ -632,7 +627,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream, is_padding);
bias, routed_scaling_factor, stream);
#else
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -640,13 +635,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream, is_padding); \
bias, routed_scaling_factor, stream); \
} else if (WARP_SIZE == 32) { \
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream, is_padding); \
bias, routed_scaling_factor, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -666,8 +661,7 @@ void topkGatingKernelLauncher(
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
cudaStream_t stream,
const bool* is_padding) {
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
@@ -742,7 +736,7 @@ void topkGatingKernelLauncher(
}
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
}
}
}
@@ -761,8 +755,7 @@ void dispatch_topk_launch(
int num_tokens, int num_experts, int topk, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
cudaStream_t stream,
std::optional<torch::stable::Tensor> is_padding)
cudaStream_t stream)
{
const float* bias_ptr = nullptr;
if (bias.has_value()) {
@@ -776,18 +769,6 @@ void dispatch_topk_launch(
bias_ptr = bias_tensor.const_data_ptr<float>();
}
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -796,7 +777,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -805,7 +786,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
} else {
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
@@ -815,7 +796,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
}
}
@@ -825,8 +806,7 @@ void topk_softmax(
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding)
std::optional<torch::stable::Tensor> bias)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -845,15 +825,15 @@ void topk_softmax(
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream, is_padding);
bias, 1.0, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream, is_padding);
bias, 1.0, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream, is_padding);
bias, 1.0, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -866,8 +846,7 @@ void topk_sigmoid(
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding)
double routed_scaling_factor)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -886,15 +865,15 @@ void topk_sigmoid(
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream, is_padding);
bias, routed_scaling_factor, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream, is_padding);
bias, routed_scaling_factor, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream, is_padding);
bias, routed_scaling_factor, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -80,27 +80,22 @@ __launch_bounds__(128) __global__
OutIndType* indices, int num_rows,
int num_experts, float routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid,
const bool* is_padding) {
const HashIndType* tid2eid) {
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane = threadIdx.x % 32;
if (warp >= num_rows) return;
const int64_t token_id = load_index_as_int64(input_ids, warp);
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
int expert = 0;
float weight = 0.f;
if (lane < 6 && !is_pad_row) {
if (lane < 6) {
// only load and calculate for 6 experts
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
const float x = input[warp * num_experts + expert];
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
if (isnan(weight)) {
weight = 0.f;
}
}
float weight_sum = weight;
#pragma unroll
@@ -116,8 +111,7 @@ __launch_bounds__(128) __global__
const int offset = warp * 6 + lane;
output[offset] =
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
: static_cast<OutIndType>(-1);
indices[offset] = static_cast<OutIndType>(expert);
}
}
@@ -126,8 +120,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
int num_rows, int num_experts,
double routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream,
const bool* is_padding) {
const HashIndType* tid2eid, cudaStream_t stream) {
if (num_rows == 0) return;
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -141,7 +134,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
config.numAttrs = 1;
const float scale = static_cast<float>(routed_scaling_factor);
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
num_experts, scale, input_ids, tid2eid, is_padding);
num_experts, scale, input_ids, tid2eid);
}
#endif
@@ -173,8 +166,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
const int num_rows, IndType* indices, int* source_rows, const int k,
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const HashIndType* input_ids, const HashIndType* tid2eid,
const bool* is_padding) {
const HashIndType* input_ids, const HashIndType* tid2eid) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -239,7 +231,6 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
@@ -258,12 +249,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
cudaGridDependencySynchronize();
#endif
if (is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = 0.f;
}
} else if constexpr (std::is_same_v<InputType, float>) {
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
// to float
if constexpr (std::is_same_v<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
const VecType* vec_thread_read_ptr =
@@ -327,22 +315,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
if constexpr (USE_HASH) {
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
row_chunk[ii] = val;
}
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
row_chunk[ii] = sqrtf(val);
}
float selected_sum = 0.f;
#pragma unroll
@@ -357,8 +335,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
if (expert == expert_idx) {
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
: static_cast<IndType>(-1);
indices[idx] = static_cast<IndType>(expert);
selected_sum += row_chunk[ii];
break;
}
@@ -402,31 +379,23 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
#endif
return;
} else {
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
// Original TopK path: find top-k experts by score
@@ -481,19 +450,18 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row =
row_is_active && node_uses_expert && !is_pad_row;
const bool should_process_row = row_is_active && node_uses_expert;
// The lead thread from each sub-group will write out the final results
// to global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
if (correction_bias != nullptr && should_process_row) {
if (correction_bias != nullptr) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] =
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -576,7 +544,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const bool use_hash, const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
const HashIndType* tid2eid, cudaStream_t stream) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
@@ -605,12 +573,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
indices, source_row, k, start_expert, end_expert,
renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, is_padding);
input_ids, tid2eid);
#else
kernel<<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert,
end_expert, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, is_padding);
input_ids, tid2eid);
#endif
})
}
@@ -624,7 +592,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
stream, is_padding);
stream);
#else
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -633,14 +601,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream, is_padding); \
tid2eid, stream); \
} else if (WARP_SIZE == 32) { \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream, is_padding); \
tid2eid, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -654,14 +622,14 @@ void topkGatingSoftplusSqrtKernelLauncher(
const int topk, const bool renormalize, double routed_scaling_factor,
const float* correction_bias, const bool use_hash,
const HashIndType* input_ids, const HashIndType* tid2eid,
cudaStream_t stream, const bool* is_padding) {
cudaStream_t stream) {
#ifndef USE_ROCM
if constexpr (std::is_same_v<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
routed_scaling_factor, input_ids, tid2eid, stream);
return;
}
}
@@ -760,8 +728,7 @@ void dispatch_topk_softplus_sqrt_launch(
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
const std::optional<torch::stable::Tensor>& is_padding) {
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -770,22 +737,6 @@ void dispatch_topk_softplus_sqrt_launch(
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
"is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
"is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (tid2eid.has_value()) {
STD_TORCH_CHECK(input_ids.has_value(),
"input_ids is required for hash MoE");
@@ -800,7 +751,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
tid2eid.value().const_data_ptr<int64_t>(), stream);
} else {
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Int);
@@ -810,7 +761,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
tid2eid.value().const_data_ptr<int>(), stream);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -818,7 +769,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
static_cast<const OutIndType*>(nullptr), stream);
}
};
@@ -842,8 +793,7 @@ void topk_softplus_sqrt(
bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding) {
const std::optional<torch::stable::Tensor>& tid2eid) {
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
@@ -856,22 +806,21 @@ void topk_softplus_sqrt(
dispatch_topk_softplus_sqrt_launch<float>(
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
token_expert_indices, num_tokens, num_experts, topk, renormalize,
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
is_padding);
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::Half) {
dispatch_topk_softplus_sqrt_launch<__half>(
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream, is_padding);
input_ids, tid2eid, stream);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream, is_padding);
input_ids, tid2eid, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
+3 -3
View File
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
"bias, Tensor? is_padding) -> ()");
"bias) -> ()");
// Apply topk sigmoid to the gating outputs.
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, "
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
"Tensor? bias, float routed_scaling_factor) -> ()");
m.def(
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, float "
"routed_scaling_factor, Tensor? "
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
// Calculate the result of moe by summing up the partial results
// from all selected experts. topk_ids/expert_map are optional and, when
-11
View File
@@ -315,17 +315,6 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
std::optional<torch::stable::Tensor> index_q_out,
const std::string& kv_cache_dtype, bool skip_index_branch);
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
void kimi_k3_attn_res(torch::stable::Tensor& prefix,
torch::stable::Tensor const& delta,
torch::stable::Tensor const& blocks,
torch::stable::Tensor const& norm_weight,
torch::stable::Tensor const& qk_weight,
torch::stable::Tensor const& output_norm_weight,
torch::stable::Tensor& output, int64_t num_blocks,
double eps, double output_norm_eps);
#endif
// Sampler kernels (shared CUDA/ROCm)
void apply_repetition_penalties_(
torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -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());
-11
View File
@@ -468,14 +468,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"int block_size, Tensor!? q_out, Tensor!? index_q_out, "
"str kv_cache_dtype, bool skip_index_branch=False) -> ()");
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
ops.def(
"kimi_k3_attn_res("
"Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
"Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
"int num_blocks, float eps, float output_norm_eps) -> ()");
#endif
// Apply repetition penalties to logits in-place.
ops.def(
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -701,9 +693,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
#endif
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
#endif
// Sampler kernels (shared CUDA/ROCm)
ops.impl("apply_repetition_penalties_",
+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)
+21 -27
View File
@@ -22,13 +22,9 @@
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
# =============================================================================
ARG CUDA_VERSION=13.0.3
ARG CUDA_VERSION=13.0.2
ARG PYTHON_VERSION=3.12
ARG UBUNTU_VERSION=22.04
# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
# This version is only used for CUDA 13+ builds; CUDA 12 falls back to
# the default NCCL version shipped with the base image.
ARG NCCL_VERSION=2.30.7
# By parameterizing the base images, we allow third-party to use their own
# base images. One use case is hermetic builds with base images stored in
@@ -298,6 +294,9 @@ FROM base AS rust-build
ARG BUILD_OS
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
# Temporary default for the initial CI validation. Set this back to 0 when
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
ARG VLLM_RUST_COVERAGE=1
# Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
@@ -481,17 +480,10 @@ WORKDIR /workspace
# Build DeepEP wheels
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
# Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
ARG DEEPEP_COMMIT_HASH=d4f41e4e93
ARG DEEPEP_COMMIT_HASH=73b6ea4
ARG NVSHMEM_VER
ARG NCCL_VERSION
RUN --mount=type=cache,target=/opt/uv/cache \
mkdir -p /tmp/ep_kernels_workspace/dist && \
CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
> /tmp/nccl-override.txt && \
export UV_OVERRIDE=/tmp/nccl-override.txt; \
fi && \
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
/tmp/install_python_libraries.sh \
--workspace /tmp/ep_kernels_workspace \
@@ -655,7 +647,6 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
ARG NCCL_VERSION
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL
@@ -708,6 +699,7 @@ RUN apt-get update -y \
# Install CUDA development tools for runtime JIT compilation
# (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \
apt-get update -y && \
apt-get install -y --no-install-recommends --allow-change-held-packages \
cuda-nvcc-${CUDA_VERSION_DASH} \
@@ -720,6 +712,12 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
libnuma-dev \
# numactl CLI for NUMA binding at runtime
numactl && \
# Fixes nccl_allocator requiring nccl.h at runtime
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
# so we pin the version to match our CUDA version
NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \
apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
rm -rf /var/lib/apt/lists/*
# Install uv for faster pip installs
@@ -739,18 +737,6 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
&& chgrp -R 0 /opt/uv \
&& chmod -R g+rwX,a+rX /opt/uv
# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins
# an older version as a transitive dep; this override forces uv to use our
# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op).
RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
if [ "$CUDA_MAJOR" -ge 13 ]; then \
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
> /etc/uv-overrides.txt; \
else \
touch /etc/uv-overrides.txt; \
fi
ENV UV_OVERRIDE=/etc/uv-overrides.txt
# ----------------------------------------------------------------------
# Non-root support (opt-in)
# ----------------------------------------------------------------------
@@ -810,7 +796,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.15.post1
ARG FLASHINFER_VERSION=0.6.14
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -919,6 +905,14 @@ COPY ./vllm/collect_env.py .
# note that this uses vllm installed by `pip`
FROM vllm-base AS test
COPY --from=rust-build \
/workspace/rust-coverage-tools/ \
/opt/vllm-rust-coverage/
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
ENV LLVM_PROFILE_FILE=/dev/null
ADD . /vllm-workspace/
ARG PYTHON_VERSION
+31 -55
View File
@@ -339,17 +339,18 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
# NIXL/UCX build stages
FROM base AS build_nixl
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
# RIXL/UCX build stages
FROM base AS build_rixl
ARG RIXL_BRANCH="39be1de8"
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
ARG UCX_BRANCH="bfb51733"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
ENV NIXL_HOME=/usr/local/nixl
ENV RIXL_HOME=/usr/local/rixl
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
# NIXL build system dependencies and RDMA support
# RIXL build system dependences and RDMA support
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
libgrpc-dev \
libgrpc++-dev \
@@ -367,8 +368,7 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
&& rm -rf /var/lib/apt/lists/*
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
uv pip install --system meson auditwheel patchelf tomlkit
RUN --mount=type=cache,target=/root/.cache/ccache \
cd /usr/local/src && \
@@ -396,50 +396,30 @@ ENV PATH=/usr/local/ucx/bin:$PATH
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
RUN --mount=type=cache,target=/root/.cache/ccache \
git clone ${NIXL_REPO} /opt/nixl && \
cd /opt/nixl && \
git checkout ${NIXL_BRANCH} && \
git clone ${RIXL_REPO} /opt/rixl && \
cd /opt/rixl && \
git checkout ${RIXL_BRANCH} && \
CC="ccache gcc" CXX="ccache g++" \
meson setup build --prefix=${NIXL_HOME} \
meson setup build --prefix=${RIXL_HOME} \
-Ducx_path=${UCX_HOME} \
-Dwheel_variant=rocm \
-Dbuild_tests=false \
-Dbuild_examples=false && \
-Drocm_path=${ROCM_PATH} && \
cd build && \
ninja -j$(nproc) && \
ninja install && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
> /etc/ld.so.conf.d/nixl.conf && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
>> /etc/ld.so.conf.d/nixl.conf && \
ldconfig
ninja install
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
# so configure the ROCm wheel variant directly through Meson.
# Generate RIXL wheel
# Exclude libcore and libpull from auditwheel: transitive dependencies
# that are not shipped in the wheel and vary across base images.
RUN cd /opt/nixl && \
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
CC="ccache gcc" CXX="ccache g++" \
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
--python ${PYTHON_VERSION} \
-Csetup-args=-Ducx_path=${UCX_HOME} \
-Csetup-args=-Dwheel_variant=rocm \
-Csetup-args=-Dbuild_tests=false \
-Csetup-args=-Dbuild_examples=false && \
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
auditwheel repair \
--exclude 'libamdhip64*' \
--exclude 'libcore*' \
--exclude 'libpull*' \
/tmp/nixl_wheels/nixl_rocm*.whl \
--plat manylinux_2_34_$(uname -m) \
--wheel-dir /tmp/nixl_wheels/repaired && \
./contrib/wheel_add_ucx_plugins.py \
RUN cd /opt/rixl && \
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
contrib/build-wheel.sh && \
mkdir -p /app/install && \
_ucx_install_dir=${UCX_HOME} \
./contrib/build-wheel.sh \
--output-dir /app/install \
--rocm-dir ${ROCM_PATH} \
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
/tmp/nixl_wheels/repaired/*.whl && \
cp /tmp/nixl_wheels/repaired/*.whl /app/install
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
# invalidate the slow ROCShmem build.
@@ -680,10 +660,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
fi
# Install NIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
# Install RIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
# Copy ROCShmem runtime libraries.
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -744,8 +724,6 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
# ROCm profiler limits workaround.
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -818,9 +796,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
&& pip uninstall -y vllm \
&& uv pip install --system *.whl
# Install NIXL ROCm wheel
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
uv pip install --system /nixl_install/*.whl
# Install RIXL wheel
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
uv pip install --system /rixl_install/*.whl
ARG COMMON_WORKDIR
ARG BASE_IMAGE
@@ -835,8 +813,6 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
ENV TOKENIZERS_PARALLELISM=false
+2 -2
View File
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
ARG FA_BRANCH="0e60e394"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="v0.1.16.post5"
ARG AITER_BRANCH="v0.1.16.post3"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
ENV AITER_ROCM_ARCH=gfx942;gfx950
# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime
ENV MORI_GPU_ARCHS=gfx942;gfx950
# Required for RCCL in ROCm7.1
ENV HSA_NO_SCRATCH_RECLAIM=1
+1 -25
View File
@@ -86,29 +86,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p /tmp/hf-xet/dist && \
cp dist/*.whl /tmp/hf-xet/dist/
# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which
# llvmlite v0.47 does not support; only SystemZ target is needed).
FROM base AS llvm20-build
ARG LLVM_VERSION=20.1.8
WORKDIR /tmp
RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \
tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \
cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
-DLLVM_ENABLE_RTTI=ON \
-DLLVM_BUILD_TOOLS=OFF \
-DLLVM_BUILD_UTILS=ON \
-DLLVM_BUILD_EXAMPLES=OFF \
-DLLVM_BUILD_TESTS=OFF \
-DLLVM_INCLUDE_TESTS=OFF \
-DLLVM_INCLUDE_EXAMPLES=OFF \
-DLLVM_INCLUDE_BENCHMARKS=OFF && \
ninja -C build install && \
rm -rf build llvm-project-${LLVM_VERSION}.src*
# Build numba
FROM python-install AS numba-builder
@@ -119,13 +96,11 @@ WORKDIR /tmp
# Clone all required dependencies
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \
microdnf install ninja-build gcc gcc-c++ -y && \
git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
cd llvmlite && \
uv pip install 'cmake<4' 'setuptools<70' numpy && \
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
python setup.py bdist_wheel && \
cd ../numba && \
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
@@ -183,6 +158,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
uv pip install -v \
$ARROW_WHL_FILE \
$VISION_WHL_FILE \
$HF_XET_WHL_FILE \
$LLVM_WHL_FILE \
+13 -13
View File
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
}
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
# stage and is irrelevant when building --target ci_base itself.
variable "CI_BASE_IMAGE" {
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
# functions produce no entries when the variable is empty.
variable "NIXL_BRANCH" {
variable "RIXL_BRANCH" {
default = ""
}
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
default = ""
}
variable "NIXL_CACHE_KEY" {
variable "RIXL_CACHE_KEY" {
default = ""
}
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
])
}
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
# These stages are pinned to specific upstream commit hashes, so cache keys use
# those hashes rather than the Buildkite commit. This means the cache persists
# across all vLLM commits as long as the upstream dependency pins don't change.
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
function "get_cache_from_rocm_deps" {
params = []
result = compact([
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
])
}
function "get_cache_to_rocm_nixl" {
function "get_cache_to_rocm_rixl" {
params = []
result = compact([
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
])
}
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
# (e.g., requirements change), these stages are cache hits if their upstream
# pins haven't changed -- saving ~35min of compilation.
target "nixl-rocm-ci" {
target "rixl-rocm-ci" {
inherits = ["_common-rocm", "_ci-rocm"]
target = "build_nixl"
target = "build_rixl"
cache-from = get_cache_from_rocm_deps()
cache-to = get_cache_to_rocm_nixl()
cache-to = get_cache_to_rocm_rixl()
output = ["type=cacheonly"]
}
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
output = ["type=cacheonly"]
}
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
# files drifts from the remote image label. Per-PR builds then pull the result
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
]),
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
# are cache hits even when ci_base itself needs rebuilding.
get_cache_from_rocm_deps(),
)
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
# Group for ci_base builds -- exports dependency stage caches alongside the
# ci_base image so future rebuilds can reuse them independently.
group "ci-base-rocm-ci-with-deps" {
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
}
+2 -2
View File
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
# overrides) and writes a small HCL override before invoking CI targets.
variable "NIXL_BRANCH" {
variable "RIXL_BRANCH" {
default = ""
}
@@ -106,7 +106,7 @@ target "test-rocm" {
output = ["type=docker"]
}
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
# the auto-rebuild trigger when requirements change in a PR.
target "ci-base-rocm" {
+8 -8
View File
@@ -2,7 +2,7 @@
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
"variable": {
"CUDA_VERSION": {
"default": "13.0.3"
"default": "13.0.2"
},
"PYTHON_VERSION": {
"default": "3.12"
@@ -10,14 +10,11 @@
"UBUNTU_VERSION": {
"default": "22.04"
},
"NCCL_VERSION": {
"default": "2.30.7"
},
"BUILD_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
},
"FINAL_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.3-base-ubuntu22.04"
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
},
"BUILD_OS": {
"default": "ubuntu"
@@ -49,6 +46,9 @@
"TORCH_CUDA_ARCH_LIST": {
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
},
"VLLM_RUST_COVERAGE": {
"default": "1"
},
"MAX_JOBS": {
"default": "2"
},
@@ -59,7 +59,7 @@
"default": "cuda"
},
"DEEPEP_COMMIT_HASH": {
"default": "d4f41e4e93"
"default": "73b6ea4"
},
"GIT_REPO_CHECK": {
"default": "0"
@@ -71,7 +71,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.15.post1"
"default": "0.6.14"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+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

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