Compare commits

..
Author SHA1 Message Date
Tyler Michael SmithandClaude 5b4f6d5284 Use cumem allocator for the KV cache by default
Enable the cumem (CUDA VMM) allocator by default on CUDA and ROCm
platforms so users get stable physical pages for KV cache without
needing to set --enable-cumem-allocator explicitly. This is required
for MNNVL KV transfers.

Changes:
- ModelConfig.enable_cumem_allocator defaults to None, resolved to
  True when the cumem C extension is available
- Cached fabric handle probe in csrc/cumem_allocator.cpp uses a real
  cuMemCreate probe instead of trusting cuDeviceGetAttribute, with
  fallback to POSIX FD when fabric handles aren't available
- Safety-net fallback: if fabric was probed as available but a real
  allocation still fails, updates the cache and retries with POSIX FD
- Weight loading bypasses cumem pool when sleep mode is off (cumem is
  only needed for KV cache stability in that case)
- Guard against stale engine allocations corrupting the singleton

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-10 15:25:35 -04:00
1388 changed files with 43329 additions and 127636 deletions
+1 -1
View File
@@ -91,7 +91,7 @@ steps:
- tests/quantization/test_cpu_wna16.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
pytest -x -v -s tests/quantization/test_cpu_wna16.py"
@@ -1,80 +0,0 @@
group: Intel
steps:
- label: ":docker: Build XPU image"
soft_fail: true
optional: true
depends_on: []
key: image-build-xpu
commands:
- bash -lc '.buildkite/image_build/image_build_xpu.sh "public.ecr.aws/q9t5s3a7" "vllm-ci-test-repo" "$BUILDKITE_COMMIT"'
env:
DOCKER_BUILDKIT: "1"
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 2
- exit_status: -10 # Agent was lost
limit: 2
- label: "XPU example Test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example'
- label: "XPU V1 test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh v1'
- label: "XPU server test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server'
@@ -5,10 +5,6 @@ steps:
- label: XPU Sleep Mode
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
-4
View File
@@ -5,10 +5,6 @@ steps:
- label: Engine (1 GPU)
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -6,10 +6,6 @@ steps:
key: eplb-algorithm
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
-4
View File
@@ -5,10 +5,6 @@ steps:
- label: vLLM IR Tests
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
-24
View File
@@ -5,10 +5,6 @@ steps:
- label: LoRA Runtime + Utils
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -38,10 +34,6 @@ steps:
- label: LoRA Fused/MoE Kernels
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -62,10 +54,6 @@ steps:
- label: LoRA Punica Kernels
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -86,10 +74,6 @@ steps:
- label: LoRA Punica FP8/XPU Ops
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -110,10 +94,6 @@ steps:
- label: LoRA Models
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -137,10 +117,6 @@ steps:
- label: LoRA Multimodal
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+2 -29
View File
@@ -5,10 +5,6 @@ steps:
- label: V1 Core + KV + Metrics
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -35,10 +31,6 @@ steps:
- label: V1 Sample + Logits
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -65,24 +57,17 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install lm_eval[api]>=0.4.12 &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
cd tests &&
pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
pytest -v -s v1/test_oracle.py &&
pytest -v -s v1/test_request.py &&
pytest -v -s v1/test_outputs.py &&
pytest -v -s v1/sample/test_topk_topp_sampler.py &&
pytest -v -s v1/sample/test_logprobs.py &&
pytest -v -s v1/sample/test_logprobs_e2e.py'
pytest -v -s v1/sample/test_topk_topp_sampler.py'
- label: XPU CPU Offload
timeout_in_minutes: 60
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -107,10 +92,6 @@ steps:
key: regression
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -142,10 +123,6 @@ steps:
timeout_in_minutes: 30
num_devices: 2
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -177,10 +154,6 @@ steps:
key: async-engine-inputs-utils-worker
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -1,62 +0,0 @@
group: Model Runner V2 Intel
depends_on:
- image-build-xpu
steps:
- label: Model Runner V2 Core Tests (Intel)
timeout_in_minutes: 45
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
- vllm/v1/core/sched/
- vllm/v1/attention/
- tests/v1/engine/test_llm_engine.py
- tests/v1/e2e/
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
cd tests &&
pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
pytest -v -s v1/e2e/general/test_min_tokens.py'
- label: Model Runner V2 Examples (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/v1/worker/gpu/
- vllm/v1/core/sched/
- vllm/v1/worker/gpu_worker.py
- examples/basic/offline_inference/
- examples/generate/multimodal/
- examples/features/
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
cd examples &&
python3 basic/offline_inference/chat.py &&
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'
@@ -6,10 +6,6 @@ steps:
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -31,10 +27,6 @@ steps:
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -55,10 +47,6 @@ steps:
key: multi-modal-models-standard-3-llava-qwen2-vl
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -80,10 +68,6 @@ steps:
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -104,10 +88,6 @@ steps:
key: multi-modal-processor
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+1 -35
View File
@@ -19,10 +19,6 @@ steps:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -53,10 +49,6 @@ steps:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -68,24 +60,19 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
- label: "XPU server test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -100,24 +87,3 @@ steps:
cd tests &&
pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
pytest -v -s benchmarks/test_serve_cli.py'
- label: "XPU quantization test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- vllm/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_auto_round.py'
@@ -6,7 +6,9 @@ tasks:
value: 0.7142
- name: "exact_match,flexible-extract"
value: 0.4579
moe_backend: "flashinfer_cutlass"
env_vars:
VLLM_USE_FLASHINFER_MOE_FP8: "1"
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
limit: 1319
num_fewshot: 5
max_model_len: 262144
@@ -68,10 +68,6 @@ def launch_lm_eval(eval_config, tp_size):
if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]:
model_args += "attention_backend=TRITON_ATTN"
moe_backend = eval_config.get("moe_backend", None)
if moe_backend is not None:
model_args += f"moe_backend={moe_backend},"
env_vars = eval_config.get("env_vars", None)
with scoped_env_vars(env_vars):
results = lm_eval.simple_evaluate(
+4 -16
View File
@@ -1,25 +1,12 @@
# CUDA architecture lists — following PyTorch RELEASE.md
# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
# These requested arches are filtered by CMake's CUDA_SUPPORTED_ARCHS before
# per-kernel arch selection. Do not add +PTX here: top-level +PTX is stripped
# during that filtering, so kernels that need PTX must request it locally.
env:
# for CUDA >=13, sm_100+ targets have family specifiers (see CMakeLists.txt)
# so targets like 10.3 and 12.1 are automatically supported with this list
CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
# aarch64-only targets: Orin (8.7), Thor (11.0, CUDA 13+)
CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0"
# for CUDA <13, we need to specify all needed targets
# some targets (10.3, 12.1) are skipped to limit the wheel size (< 500MB)
# please use CUDA 13 wheels or compile yourself on these new devices
CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
# aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
# pre-built mooncake wheels
# the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04
# so we use different wheels for the time being
MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"
@@ -859,6 +846,7 @@ steps:
allow_failure: true
- step: build-cpu-release-image-arm64
allow_failure: true
if: build.env("NIGHTLY") != "1"
- label: "Publish release images to DockerHub"
depends_on:
-3
View File
@@ -13,8 +13,5 @@ INPUT_FILE="$1"
# Strip timestamps
sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE"
# Strip Buildkite inline timestamp markers (ESC _bk;t=<ms> BEL)
sed -i 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE"
# Strip colorization
sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE"
+47 -151
View File
@@ -1,178 +1,74 @@
#!/bin/bash
# Fetch vLLM Buildkite CI logs (public; no login required).
# Usage: ./ci-fetch-log.sh <buildkite_job_url> [output_file]
# ./ci-fetch-log.sh <build_number> <job_uuid> [output_file]
#
# Usage:
# ci-fetch-log.sh [--soft|--all] --pr [<PR>] failed jobs in the PR's latest
# build (current branch if omitted)
# ci-fetch-log.sh [--soft|--all] <build_url> failed jobs in that build
# ci-fetch-log.sh <job_url> [output] one job; both #<job_uuid> and
# ?sid=<id> URL forms work
# ci-fetch-log.sh <build> <job_uuid> [output]
# Downloads the raw log for a Buildkite job from the public, unauthenticated
# /organizations/<org>/pipelines/<pipeline>/builds/<n>/jobs/<uuid>/download
# endpoint, then strips ANSI/timestamps via ci-clean-log.sh.
#
# --soft also fetches soft-failed jobs; --all fetches every finished job.
# Saves each log as ci-<build>-<job-name>.log (ANSI/timestamps stripped) and
# prints "<file>\t<job name>" per job. [output] is single-job only; "-"
# streams to stdout. Existing files are kept; CI_FETCH_LOG_FORCE=1 refetches.
# Find <build_number> and <job_uuid> via:
# gh pr checks <PR> --repo vllm-project/vllm
# Each failing row's URL is .../builds/<build_number>#<job_uuid>.
#
# Default output path: ci-<build>-<uuid_first_13_chars>.log (e.g.
# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's
# first 8 chars, so the second segment is needed for uniqueness when
# fetching multiple jobs in parallel. The script refuses to overwrite an
# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1
# to override.
set -euo pipefail
ORG="vllm"
PIPELINE="ci"
UA="vllm-ci-fetch-log"
UUID_RE='[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}'
usage() {
sed -n '2,15p' "$0" | sed 's/^# \{0,1\}//'
echo "Usage: $0 <buildkite_job_url> [output_file]"
echo " $0 <build_number> <job_uuid> [output_file]"
exit 1
}
die() {
echo "$1" >&2
exit 1
}
if [ $# -lt 1 ]; then usage; fi
BUILD="" JOB="" SID="" OUT=""
SCOPE="failed"
while :; do
case "${1:-}" in
--soft) SCOPE="soft" ;;
--all) SCOPE="all" ;;
*) break ;;
esac
shift
done
case "${1:-}" in
--pr)
PR="${2:-}"
# gh pr checks exits non-zero when checks are failing; that is the
# expected case here.
URL=$(gh pr checks ${PR:+"$PR"} --repo vllm-project/vllm 2>/dev/null |
grep -oE "https://buildkite.com/${ORG}/${PIPELINE}/builds/[0-9]+" |
sort -t/ -k7 -n | tail -1 || true)
[ -n "$URL" ] || die "No Buildkite build found via: gh pr checks ${PR:-<current branch>}"
BUILD="${URL##*/}"
;;
https://*)
if [[ "$1" == https://* ]]; then
BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
JOB=$(echo "$1" | grep -oE "#${UUID_RE}" | head -n 1 | cut -c2- || true)
SID=$(echo "$1" | grep -oE "[?&]sid=${UUID_RE}" | head -n 1 | sed 's/.*sid=//' || true)
JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
OUT="${2:-}"
[ -n "$BUILD" ] || die "Could not parse build number from: $1"
;;
[0-9]*)
[ $# -ge 2 ] || usage
else
if [ $# -lt 2 ]; then usage; fi
BUILD="$1"
JOB="$2"
OUT="${3:-}"
;;
*)
fi
if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
echo "Could not parse build number or job UUID from: $1" >&2
usage
;;
esac
fi
# Jobs in the same build share the UUID's first segment, so include the
# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames
# unique when fetching multiple jobs from one build in parallel.
if [ -z "$OUT" ]; then
OUT="ci-${BUILD}-${JOB:0:13}.log"
fi
if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2
exit 1
fi
COOKIES=$(mktemp)
JOBS_TSV=$(mktemp)
trap 'rm -f "$COOKIES" "$JOBS_TSV"' EXIT
trap 'rm -f "$COOKIES"' EXIT
# Buildkite issues a session cookie on first hit; later requests need it.
curl -fsSL -c "$COOKIES" -A "$UA" \
# Buildkite issues a session cookie on first hit; subsequent /download needs it.
curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \
"https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null
# The build's job list (id, step uuid, state, name) is served as JSON from
# the user-facing /data/jobs endpoint. Flatten it to TSV for easy filtering:
# job_id step_uuid failed soft_failed finished slug name
curl -fsSL -b "$COOKIES" -A "$UA" \
"https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}/data/jobs" |
python3 -c '
import json, re, sys
curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \
"https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \
-o "$OUT"
data = json.load(sys.stdin)
if data.get("has_next_page"):
print("warning: job list is paginated; some jobs not shown", file=sys.stderr)
for r in data["records"]:
if r.get("type") != "script":
continue
name = (r.get("name") or "").replace("\t", " ").replace("\n", " ")
slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:60]
print("\t".join([
r["id"],
r.get("step_uuid") or "",
str(r.get("passed") is False),
str(bool(r.get("soft_failed"))),
str(bool(r.get("finished_at"))),
slug,
name,
]))
' >"$JOBS_TSV" || die "Could not list jobs for build ${BUILD}"
bash "$(dirname "$0")/ci-clean-log.sh" "$OUT"
if [ -n "$SID" ] && [ -z "$JOB" ]; then
# The ?sid= in builds/<N>/list URLs is the *step* uuid, not the job uuid.
JOB=$(awk -F'\t' -v s="$SID" '$1 == s || $2 == s {print $1; exit}' "$JOBS_TSV")
[ -n "$JOB" ] || die "No job matching sid=${SID} in build ${BUILD}"
fi
fetch_job() { # <job_uuid> <output_file>
curl -fsSL -b "$COOKIES" -A "$UA" \
"https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/$1/download" \
-o "$2"
bash "$(dirname "$0")/ci-clean-log.sh" "$2"
}
if [ -n "$JOB" ]; then
# Single-job mode.
NAME=$(awk -F'\t' -v j="$JOB" '$1 == j {print $7; exit}' "$JOBS_TSV")
SLUG=$(awk -F'\t' -v j="$JOB" '$1 == j {print $6; exit}' "$JOBS_TSV")
[ -n "$OUT" ] || OUT="ci-${BUILD}-${SLUG:-${JOB:0:13}}.log"
if [ "$OUT" = "-" ]; then
TMP=$(mktemp)
fetch_job "$JOB" "$TMP"
cat "$TMP"
rm -f "$TMP"
exit 0
fi
if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
die "Refusing to overwrite existing ${OUT} (set CI_FETCH_LOG_FORCE=1 or pass an output path)."
fi
fetch_job "$JOB" "$OUT"
printf '%s\t%s\n' "$OUT" "${NAME:-$JOB}"
exit 0
fi
# Build-wide mode: fetch finished jobs matching $SCOPE.
[ -z "$OUT" ] || die "[output_file] is only valid when fetching a single job."
case "$SCOPE" in
failed) FILTER='$3 == "True" && $4 == "False" && $5 == "True"' ;;
soft) FILTER='$3 == "True" && $5 == "True"' ;;
all) FILTER='$5 == "True"' ;;
esac
if [ "$SCOPE" = "failed" ]; then
SOFT=$(awk -F'\t' '$3 == "True" && $4 == "True"' "$JOBS_TSV" | wc -l)
[ "$SOFT" -eq 0 ] || echo "Skipping ${SOFT} soft-failed job(s); use --soft to include them." >&2
fi
FOUND=0
EMITTED=" "
while IFS=$'\t' read -r job_id _ _ _ _ slug name; do
FOUND=$((FOUND + 1))
out="ci-${BUILD}-${slug:-${job_id:0:13}}.log"
# Retries share a name with the original job; disambiguate by uuid.
case "$EMITTED" in
*" $out "*) out="ci-${BUILD}-${slug:-job}-${job_id:0:13}.log" ;;
esac
EMITTED="${EMITTED}${out} "
if [ -e "$out" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
echo "Keeping existing ${out} (set CI_FETCH_LOG_FORCE=1 to refetch)." >&2
elif ! fetch_job "$job_id" "$out"; then
echo "Failed to download log for job ${job_id} (${name})." >&2
continue
fi
printf '%s\t%s\n' "$out" "$name"
done < <(awk -F'\t' "$FILTER" "$JOBS_TSV")
if [ "$FOUND" -eq 0 ]; then
echo "No matching jobs in build ${BUILD} (scope: ${SCOPE})." >&2
fi
echo "$OUT"
@@ -1,51 +0,0 @@
#!/bin/bash
set -euo pipefail
test_suite="${1:-}"
if [[ -z "${test_suite}" ]]; then
echo "Usage: $0 <example|v1|server>" >&2
exit 1
fi
case "${test_suite}" in
example)
pip install tblib==3.1.0
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8
python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
;;
v1)
cd tests
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp"
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
pytest -v -s v1/structured_output
pytest -v -s v1/test_serial_utils.py
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
;;
server)
pip install av
cd tests
pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py
pytest -v -s benchmarks/test_serve_cli.py
;;
*)
echo "Unknown Intel test suite: ${test_suite}" >&2
exit 1
;;
esac
@@ -243,10 +243,8 @@ container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head
# ---- Command source selection ----
commands=""
commands_source=""
if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
commands="${VLLM_TEST_COMMANDS}"
commands_source="env"
echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
elif [[ $# -gt 0 ]]; then
all_yaml=true
@@ -305,12 +303,8 @@ if [[ -z "$commands" ]]; then
fi
echo "Raw commands: $commands"
if [[ "$commands_source" != "env" ]]; then
commands=$(re_quote_pytest_markers "$commands")
echo "After re-quoting: $commands"
else
echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
fi
commands=$(re_quote_pytest_markers "$commands")
echo "After re-quoting: $commands"
commands=$(apply_intel_test_overrides "$commands")
echo "Final commands: $commands"
@@ -4,11 +4,6 @@
set -euo pipefail
if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then
uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
exit 0
fi
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
uv pip install --system -r "${REQUIREMENTS_FILE}"
@@ -110,36 +110,6 @@ install_uv() {
| env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
}
setup_pyo3_python() {
local python_version="${PYO3_PYTHON_VERSION:-3.12}"
log_section "Installing Python ${python_version} for PyO3 tests"
uv python install "$python_version"
PYO3_PYTHON="$(uv python find \
--managed-python \
--no-project \
--resolve-links \
"$python_version")"
export PYO3_PYTHON
local python_libdir
python_libdir="$("$PYO3_PYTHON" - <<'PY'
import pathlib
import sysconfig
libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
ldlibrary = sysconfig.get_config_var("LDLIBRARY")
assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
assert ldlibrary
assert (libdir / ldlibrary).exists(), libdir / ldlibrary
print(libdir)
PY
)"
export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
}
run_style_clippy() {
install_cargo_sort
@@ -162,7 +132,6 @@ run_style_clippy() {
run_tests() {
install_uv
setup_pyo3_python
install_cargo_nextest
log_section "Running cargo nextest"
+50 -38
View File
@@ -398,7 +398,7 @@ steps:
- tests/kernels/helion/
- vllm/platforms/rocm.py
commands:
- pip install helion==1.1.0
- pip install helion==1.0.0
- pytest -v -s kernels/helion/
- label: Kernels Mamba Test # TBD
@@ -415,6 +415,22 @@ steps:
commands:
- pytest -v -s kernels/mamba
#----------------------------------------------------------- mi250 · lora ------------------------------------------------------------#
- label: LoRA %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
#------------------------------------------------------ mi250 · models / basic -------------------------------------------------------#
- label: Basic Models Test (Other CPU) # TBD
@@ -592,11 +608,6 @@ steps:
- pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
- pip uninstall bge_m3_sparse_plugin -y
# END: `bge_m3_sparse io_processor` test
# BEGIN: `colbert_query io_processor` test
- pip install -e ./plugins/colbert_query_plugin
- pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
- pip uninstall colbert_query_plugin -y
# END: `colbert_query io_processor` test
# BEGIN: `stat_logger` plugins test
- pip install -e ./plugins/vllm_add_dummy_stat_logger
- pytest -v -s plugins_tests/test_stats_logger_plugins.py
@@ -647,7 +658,7 @@ steps:
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 35
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
@@ -1594,10 +1605,9 @@ steps:
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 55
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
parallelism: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -1628,11 +1638,10 @@ steps:
- 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 MoE Test %N # TBD
timeout_in_minutes: 50
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
parallelism: 5
parallelism: 4
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
@@ -1685,20 +1694,6 @@ steps:
#----------------------------------------------------------- mi300 · lora ------------------------------------------------------------#
- label: LoRA %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
- label: LoRA TP (Distributed) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1783,9 +1778,10 @@ steps:
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
- pytest -v -s models/multimodal/test_mapping.py
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-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 hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- label: Multi-Modal Models (Extended Generation 2) # TBD
timeout_in_minutes: 180
@@ -1797,8 +1793,9 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-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)'
- label: Multi-Modal Models (Extended Generation 3) # TBD
@@ -2077,6 +2074,19 @@ steps:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
- label: e2e Scheduling (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2122,10 +2132,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: Spec Decode Eagle # TBD
timeout_in_minutes: 45
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
@@ -2745,7 +2754,7 @@ steps:
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k"
- python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
#-------------------------------------------------------- mi355 · distributed --------------------------------------------------------#
@@ -2937,6 +2946,7 @@ steps:
- vllm/model_executor/models/qwen3_5_mtp.py
- vllm/transformers_utils/configs/qwen3_5.py
- vllm/transformers_utils/configs/qwen3_5_moe.py
- vllm/model_executor/models/qwen.py
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_next.py
@@ -3043,7 +3053,7 @@ steps:
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
- label: Kernels (B200-MI355) # TBD
timeout_in_minutes: 15
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
working_dir: "/vllm-workspace/"
@@ -3067,10 +3077,11 @@ steps:
- pytest -v -s tests/kernels/attention/test_attention_selector.py
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 60
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
parallelism: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/attention/
@@ -3084,10 +3095,10 @@ steps:
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 50
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
parallelism: 5
parallelism: 4
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
@@ -3173,6 +3184,7 @@ steps:
- vllm/model_executor/models/qwen3_5_mtp.py
- vllm/transformers_utils/configs/qwen3_5.py
- vllm/transformers_utils/configs/qwen3_5_moe.py
- vllm/model_executor/models/qwen.py
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_next.py
+1 -1
View File
@@ -15,7 +15,7 @@ steps:
- pytest -v -s v1/attention
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 70
depends_on:
- image-build-amd
+4 -3
View File
@@ -12,6 +12,7 @@ steps:
- tests/basic_correctness/test_cpu_offload
- tests/basic_correctness/test_mem.py
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_mem.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_basic_correctness.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_cpu_offload.py
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s basic_correctness/test_mem.py
- pytest -v -s basic_correctness/test_basic_correctness.py
- pytest -v -s basic_correctness/test_cpu_offload.py
+1 -1
View File
@@ -23,4 +23,4 @@ steps:
- benchmarks/attention_benchmarks/
- vllm/v1/attention/
commands:
- python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k"
- python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
+10 -5
View File
@@ -14,7 +14,8 @@ steps:
- vllm/v1/cudagraph_dispatcher.py
- tests/compile/correctness_e2e/test_sequence_parallel.py
commands:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- label: Sequence Parallel Correctness Tests (2xH100)
key: sequence-parallel-correctness-tests-2xh100
@@ -24,7 +25,8 @@ steps:
optional: true
num_devices: 2
commands:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- label: AsyncTP Correctness Tests (2xH100)
key: asynctp-correctness-tests-2xh100
@@ -34,7 +36,8 @@ steps:
optional: true
num_devices: 2
commands:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- label: AsyncTP Correctness Tests (B200)
key: asynctp-correctness-tests-b200
@@ -44,7 +47,8 @@ steps:
optional: true
num_devices: 2
commands:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- label: Distributed Compile Unit Tests (2xH100)
key: distributed-compile-unit-tests-2xh100
@@ -57,7 +61,8 @@ steps:
- vllm/model_executor/layers
- tests/compile/passes/distributed/
commands:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -s -v tests/compile/passes/distributed
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- pytest -s -v tests/compile/passes/distributed
- label: Fusion and Compile Unit Tests (2xB200)
key: fusion-and-compile-unit-tests-2xb200
-14
View File
@@ -61,20 +61,6 @@ steps:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_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
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- vllm/v1/core/kv_cache_coordinator.py
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
timeout_in_minutes: 30
+46 -37
View File
@@ -32,10 +32,11 @@ steps:
- tests/entrypoints/openai/test_multi_api_servers.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
- label: Distributed Compile + RPC Tests (2 GPUs)
key: distributed-compile-rpc-tests-2-gpus
@@ -55,9 +56,10 @@ steps:
- tests/entrypoints/llm/test_collective_rpc.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s entrypoints/llm/test_collective_rpc.py
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/fullgraph/test_basic_correctness.py
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/test_wrapper.py
- export NCCL_CUMEM_HOST_ENABLE=0
- 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
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
key: distributed-torchrun-shutdown-tests-2-gpus
@@ -76,10 +78,11 @@ steps:
- tests/v1/worker/test_worker_memory_snapshot.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- NCCL_CUMEM_HOST_ENABLE=0 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'
- NCCL_CUMEM_HOST_ENABLE=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/worker/test_worker_memory_snapshot.py
- export NCCL_CUMEM_HOST_ENABLE=0
- 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 Torchrun + Examples (4 GPUs)
key: distributed-torchrun-examples-4-gpus
@@ -94,23 +97,24 @@ steps:
- tests/examples/features/data_parallel/data_parallel_offline.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
# test with torchrun tp=2 and external_dp=2
- NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
- torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# test with torchrun tp=2 and pp=2
- NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
- PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# test with torchrun tp=4 and dp=1
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with torchrun tp=2, pp=2 and dp=1
- NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with torchrun tp=1 and dp=4 with ep
- NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with torchrun tp=2 and dp=2 with ep
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with internal dp
- NCCL_CUMEM_HOST_ENABLE=0 python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
# rlhf examples
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
- label: Distributed DP Tests (4 GPUs)
key: distributed-dp-tests-4-gpus
@@ -124,13 +128,14 @@ steps:
- tests/distributed/test_utils
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_utils.py
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
- pytest -v -s distributed/test_utils.py
- label: Distributed Compile + Comm (4 GPUs)
key: distributed-compile-comm-4-gpus
@@ -146,12 +151,13 @@ steps:
- tests/distributed/test_multiproc_executor.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s compile/fullgraph/test_basic_correctness.py
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pynccl.py
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_events.py
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_symm_mem_allreduce.py
- export NCCL_CUMEM_HOST_ENABLE=0
- 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
# test multi-node TP with multiproc executor (simulated on single node)
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
- label: Distributed Tests (8 GPUs)(H100)
key: distributed-tests-8-gpus-h100
@@ -170,8 +176,9 @@ steps:
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
# test with torchrun tp=2 and dp=4 with ep
- NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- 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 Tests (4 GPUs)(A100)
key: distributed-tests-4-gpus-a100
@@ -264,7 +271,9 @@ steps:
- tests/distributed/test_pipeline_parallel.py
- tests/basic_correctness/test_basic_correctness.py
commands:
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor.py
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor_e2e.py
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
- export NCCL_CUMEM_HOST_ENABLE=0
- pytest -v -s distributed/test_ray_v2_executor.py
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
+2 -12
View File
@@ -28,7 +28,7 @@ steps:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
@@ -44,7 +44,7 @@ steps:
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 40
depends_on:
- image-build-amd
@@ -74,16 +74,6 @@ steps:
- tests/v1/e2e/general/
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
mirror:
amd:
device: mi250_1
timeout_in_minutes: 35
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
- label: V1 e2e (2 GPUs)
key: v1-e2e-2-gpus
+23 -16
View File
@@ -22,9 +22,10 @@ steps:
- vllm/
- tests/entrypoints/llm
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
mirror:
amd:
device: mi325_1
@@ -40,11 +41,12 @@ steps:
- vllm/
- tests/entrypoints/serve
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
mirror:
amd:
device: mi325_1
device: mi300_1
depends_on:
- image-build-amd
@@ -57,10 +59,11 @@ steps:
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 80
depends_on:
- image-build-amd
@@ -74,11 +77,12 @@ steps:
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 80
depends_on:
- image-build-amd
@@ -100,7 +104,7 @@ steps:
- pytest -v -s entrypoints/anthropic
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
@@ -124,7 +128,8 @@ steps:
- vllm/
- tests/entrypoints/speech_to_text
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/speech_to_text
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (Multimodal)
device: h200_35gb
@@ -135,7 +140,8 @@ steps:
- vllm/
- tests/entrypoints/multimodal
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/multimodal
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/multimodal
- label: Entrypoints Integration (Pooling)
key: entrypoints-integration-pooling
@@ -145,7 +151,8 @@ steps:
- vllm/
- tests/entrypoints/pooling
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/pooling
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/pooling
- label: OpenAI API Correctness
key: openai-api-correctness
@@ -158,7 +165,7 @@ steps:
- pytest -s entrypoints/openai/correctness/
mirror:
amd:
device: mi325_1
device: mi300_1
depends_on:
- image-build-amd
source_file_dependencies:
+2 -46
View File
@@ -74,33 +74,6 @@ steps:
commands:
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 2
mirror:
amd:
device: mi325_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/attention/
- vllm/v1/attention
- vllm/model_executor/layers/attention
- tests/kernels/attention
- vllm/_aiter_ops.py
- vllm/envs.py
- vllm/platforms/rocm.py
- label: Kernels Attention DiffKV Test (H100)
key: kernels-attention-diffkv-test-h100
timeout_in_minutes: 20
device: h100
num_devices: 1
source_file_dependencies:
- vllm/v1/attention/ops/triton_unified_attention_diffkv.py
- vllm/v1/attention/backends/triton_attn_diffkv.py
- vllm/v1/attention/backends/flash_attn_diffkv.py
- tests/kernels/attention/test_triton_unified_attention_diffkv.py
commands:
- pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
- label: Kernels Quantization Test %N
key: kernels-quantization-test
@@ -114,11 +87,10 @@ steps:
parallelism: 2
mirror:
amd:
device: mi325_1
device: mi300_1
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
- vllm/config/
- tests/kernels/quantization
- tests/kernels/quantization/test_rocm_skinny_gemms.py
- vllm/_aiter_ops.py
@@ -142,22 +114,6 @@ steps:
- pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 5
mirror:
amd:
device: mi325_1
timeout_in_minutes: 50
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
- tests/kernels/moe
- vllm/model_executor/layers/fused_moe/
- vllm/distributed/device_communicators/
- vllm/envs.py
- vllm/config
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
depends_on:
- image-build-amd
- label: Kernels Mamba Test
key: kernels-mamba-test
@@ -268,7 +224,7 @@ steps:
- vllm/utils/import_utils.py
- tests/kernels/helion/
commands:
- pip install helion==1.1.0
- pip install helion==1.0.0
- pytest -v -s kernels/helion/
+3 -10
View File
@@ -14,7 +14,7 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
@@ -50,7 +50,8 @@ steps:
- csrc/
- vllm/model_executor/layers/quantization
commands:
- VLLM_USE_DEEP_GEMM=0 pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4 # Triton is faster than DeepGEMM for H100
- export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
- label: LM Eval Small Models (B200)
key: lm-eval-small-models-b200
@@ -100,14 +101,6 @@ steps:
num_devices: 8
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
mirror:
amd:
device: mi300_8
timeout_in_minutes: 180
depends_on:
- image-build-amd
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_ROCM_ARCH=gfx942 pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt # Limit Quark compilation to save time
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
key: moe-refactor-integration-test-h100-temporary
+8 -19
View File
@@ -12,17 +12,6 @@ steps:
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
parallelism: 4
mirror:
amd:
device: mi325_1
working_dir: "/vllm-workspace/tests"
timeout_in_minutes: 60
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
depends_on:
- image-build-amd
- label: LoRA TP (Distributed)
@@ -36,14 +25,14 @@ steps:
commands:
# FIXIT: find out which code initialize cuda before running the test
# before the fix, we need to use spawn to test it
#
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# Alot of these tests are on the edge of OOMing
#
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# There is some Tensor Parallelism related processing logic in LoRA that
# requires multi-GPU testing for validation.
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_chatglm3_tp.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_llama_tp.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen3_with_multi_loras.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_olmoe_tp.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_gptoss_tp.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen35_densemodel_lora.py
- pytest -v -s -x lora/test_chatglm3_tp.py
- pytest -v -s -x lora/test_llama_tp.py
- 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
+38 -58
View File
@@ -18,14 +18,9 @@ steps:
- vllm/v1/
- tests/v1/spec_decode
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# TODO: create another `optional` test group for slow tests
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not slow_test' v1/spec_decode
mirror:
amd:
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
- pytest -v -s -m 'not slow_test' v1/spec_decode
- label: V1 Sample + Logits
key: v1-sample-logits
@@ -49,11 +44,12 @@ steps:
- tests/v1/test_request.py
- tests/v1/test_outputs.py
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/logits_processors
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_oracle.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_request.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_outputs.py
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s v1/sample
- pytest -v -s v1/logits_processors
- pytest -v -s v1/test_oracle.py
- pytest -v -s v1/test_request.py
- pytest -v -s v1/test_outputs.py
mirror:
amd:
device: mi325_1
@@ -91,23 +87,18 @@ steps:
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# split the test to avoid interference
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/core
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/executor
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_offload
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/simple_kv_offload
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/worker
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/metrics
- pytest -v -s -m 'not cpu_test' v1/core
- pytest -v -s v1/executor
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
- pytest -v -s v1/worker
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
# Integration test for streaming correctness (requires special branch).
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
mirror:
amd:
device: mi325_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
- label: V1 Others (CPU)
key: v1-others-cpu
@@ -147,23 +138,10 @@ steps:
- vllm/v1/spec_decode/extract_hidden_states.py
- vllm/model_executor/models/extract_hidden_states.py
- vllm/transformers_utils/configs/extract_hidden_states.py
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
- tests/v1/kv_connector/extract_hidden_states_integration
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_connector/extract_hidden_states_integration
- label: Extract Hidden States Integration (2 GPUs)
key: extract-hidden-states-integration-2-gpus
timeout_in_minutes: 20
num_devices: 2
source_file_dependencies:
- vllm/v1/spec_decode/extract_hidden_states.py
- vllm/model_executor/models/extract_hidden_states.py
- vllm/transformers_utils/configs/extract_hidden_states.py
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
- tests/v1/kv_connector/extract_hidden_states_integration
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
- label: Regression
key: regression
@@ -315,7 +293,6 @@ steps:
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/test_envs.py
- tests/test_inputs.py
- tests/test_outputs.py
- tests/test_pooling_params.py
@@ -323,25 +300,24 @@ steps:
- tests/multimodal
- tests/renderers
- tests/standalone_tests/lazy_imports.py
- tests/tokenizers_
- tests/reasoning
- tests/tool_parsers
- tests/tokenizers_
- tests/parser
- tests/transformers_utils
- tests/config
device: cpu-small
commands:
- python3 standalone_tests/lazy_imports.py
- pytest -v -s test_envs.py
- pytest -v -s test_inputs.py
- pytest -v -s test_outputs.py
- pytest -v -s test_pooling_params.py
- pytest -v -s test_ray_env.py
- pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s renderers
- pytest -v -s tokenizers_
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
- pytest -v -s tool_parsers
- pytest -v -s tokenizers_
- pytest -v -s parser
- pytest -v -s transformers_utils
- pytest -v -s config
@@ -355,9 +331,10 @@ steps:
- vllm/model_executor/layers
- tests/v1/determinism/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn 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]
- 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[TRITON_MLA]
- label: Batch Invariance (H100)
key: batch-invariance-h100
@@ -368,11 +345,12 @@ steps:
- vllm/model_executor/layers
- tests/v1/determinism/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn 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_WORKER_MULTIPROC_METHOD=spawn 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_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[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
@@ -383,13 +361,14 @@ steps:
- vllm/model_executor/layers
- tests/v1/determinism/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn 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_WORKER_MULTIPROC_METHOD=spawn 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]
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
- 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[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
- label: Acceptance Length Test (Large Models) # optional
device: h200_35gb
@@ -404,4 +383,5 @@ steps:
- vllm/model_executor/models/mlp_speculator.py
- tests/v1/spec_decode/test_acceptance_length.py
commands:
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
+6 -5
View File
@@ -13,12 +13,13 @@ steps:
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
# init surfaces a stack instead of silently stalling.
- export PYTHONFAULTHANDLER=1
# Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
# init) fails fast with a traceback instead of running until the global
# build timeout. The `thread` method also handles hangs inside C/CUDA
# calls that the signal method cannot interrupt.
#
# Env vars are inlined because CONTINUE_ON_FAILURE wraps each command
# in a subshell, so a standalone `export` would be lost.
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
- pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
+35 -29
View File
@@ -16,14 +16,15 @@ steps:
- tests/entrypoints/llm/test_struct_output_generate.py
commands:
- set -x
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
# This requires eager until we sort out CG correctness issues.
# TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
- VLLM_USE_V2_MODEL_RUNNER=1 ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_context_length.py
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_min_tokens.py
- 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
# Temporary hack filter to exclude ngram spec decoding based tests.
- VLLM_USE_V2_MODEL_RUNNER=1 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 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"
- label: Model Runner V2 Examples
device: h200_35gb
@@ -41,25 +42,26 @@ steps:
- examples/features/tensorize_vllm_model.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pip install tensorizer # for tensorizer test
- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/chat.py # for basic
- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model facebook/opt-125m
#- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
#- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/embed.py # TODO
- 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
#- python3 basic/offline_inference/embed.py # TODO
# for multi-modal models
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/audio_language_offline.py --seed 0
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_offline.py --seed 0
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
# for pooling models
- VLLM_USE_V2_MODEL_RUNNER=1 python3 pooling/embed/vision_embedding_offline.py --seed 0
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# for features demo
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/automatic_prefix_caching/prefix_caching_offline.py
- VLLM_USE_V2_MODEL_RUNNER=1 python3 deployment/llm_engine_example.py
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 deployment/llm_engine_example.py
- python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
- VLLM_USE_V2_MODEL_RUNNER=1 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
- 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
- label: Model Runner V2 Distributed (2 GPUs)
key: model-runner-v2-distributed-2-gpus
@@ -74,11 +76,13 @@ steps:
- tests/v1/distributed/test_eagle_dp.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
# The "and not True" here is a hacky way to exclude the prompt_embeds cases which aren't yet supported.
- VLLM_USE_V2_MODEL_RUNNER=1 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
# https://github.com/NVIDIA/nccl/issues/1838
- VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
- VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
key: model-runner-v2-pipeline-parallelism-4-gpus
@@ -93,9 +97,10 @@ steps:
- tests/v1/distributed/test_pp_dp_v2.py
commands:
- set -x
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/distributed/test_pp_dp_v2.py
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
- pytest -v -s v1/distributed/test_pp_dp_v2.py
- label: Model Runner V2 Spec Decode
device: h200_35gb
@@ -110,7 +115,8 @@ steps:
- tests/v1/e2e/spec_decode/test_spec_decode.py
commands:
- set -x
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
+6 -11
View File
@@ -15,7 +15,7 @@ steps:
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
mirror:
amd:
device: mi325_1
device: mi300_1
depends_on:
- image-build-amd
@@ -30,9 +30,10 @@ steps:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
mirror:
amd:
device: mi325_1
device: mi300_1
depends_on:
- image-build-amd
@@ -49,7 +50,7 @@ steps:
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
mirror:
amd:
device: mi325_1
device: mi300_1
depends_on:
- image-build-amd
@@ -62,15 +63,9 @@ steps:
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
mirror:
amd:
device: mi325_1
depends_on:
- image-build-amd
- label: Multi-Modal Processor (CPU)
key: multi-modal-processor-cpu
@@ -160,7 +155,7 @@ steps:
- pytest -v -s models/multimodal/pooling -m 'not core_model'
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
-18
View File
@@ -27,10 +27,6 @@ steps:
- pip install -e ./plugins/bge_m3_sparse_plugin
- pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
- pip uninstall bge_m3_sparse_plugin -y
# test colbert_query io_processor plugin
- pip install -e ./plugins/colbert_query_plugin
- pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
- pip uninstall colbert_query_plugin -y
# end io_processor plugins test
# begin stat_logger plugins test
- pip install -e ./plugins/vllm_add_dummy_stat_logger
@@ -44,17 +40,3 @@ steps:
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
- label: GGUF Plugin
key: gguf-plugin
device: h200_18gb
timeout_in_minutes: 30
soft_fail: true
optional: true
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/plugins_tests/test_gguf_plugin.py
commands:
- pip install "vllm-gguf-plugin >= 0.0.2"
- pytest -v -s plugins_tests/gguf
-6
View File
@@ -107,12 +107,6 @@ steps:
- tests/compile/passes
commands:
- pytest -s -v compile/passes --ignore compile/passes/distributed
mirror:
amd:
device: mi300_1
timeout_in_minutes: 180
depends_on:
- image-build-amd
- label: PyTorch Fullgraph Smoke Test
key: pytorch-fullgraph-smoke-test
-12
View File
@@ -21,18 +21,6 @@ steps:
- uv pip install --system conch-triton-kernels
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
- label: Quantized Fusions
key: quantized-fusions
timeout_in_minutes: 30
source_file_dependencies:
- tests/fusion
- vllm/model_executor/layers/fusion
- vllm/model_executor/kernels/linear
- vllm/model_executor/layers/quantization/compressed_tensors
- vllm/model_executor/layers/quantization/modelopt.py
commands:
- pytest -v -s fusion/
- label: Quantized MoE Test (B200)
key: quantized-moe-test-b200
timeout_in_minutes: 60
+27 -20
View File
@@ -23,15 +23,17 @@ steps:
# - tests/entrypoints/openai/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
commands:
- VLLM_USE_RUST_FRONTEND=1 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)"
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
# - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_return_token_ids.py
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_uds.py
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
# - pytest -v -s entrypoints/openai/test_return_token_ids.py
# - pytest -v -s entrypoints/openai/test_uds.py
- pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
- label: Rust Frontend Serve/Admin Coverage
timeout_in_minutes: 60
@@ -49,11 +51,13 @@ steps:
- tests/entrypoints/serve/instrumentator/test_metrics.py
# - tests/entrypoints/serve/dev/test_sleep.py
commands:
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/test_sleep.py
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
- label: Rust Frontend Core Correctness
timeout_in_minutes: 30
@@ -65,7 +69,9 @@ steps:
- tests/utils.py
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
- 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
- label: Rust Frontend Tool Use
timeout_in_minutes: 60
@@ -77,7 +83,9 @@ steps:
- tests/utils.py
- tests/tool_use/
commands:
- VLLM_USE_RUST_FRONTEND=1 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"
- 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"
- label: Rust Frontend Distributed
timeout_in_minutes: 30
@@ -91,10 +99,9 @@ steps:
- vllm/v1/engine/
- vllm/v1/worker/
- tests/utils.py
- tests/v1/distributed/test_external_lb_dp.py
- tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py
commands:
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
+5 -18
View File
@@ -12,20 +12,6 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
mirror:
amd:
device: mi325_1
timeout_in_minutes: 45
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode Eagle Nightly B200
key: spec-decode-eagle-nightly-b200
@@ -53,7 +39,7 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
@@ -92,7 +78,7 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
@@ -117,7 +103,7 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
mirror:
amd:
device: mi325_1
device: mi300_1
timeout_in_minutes: 50
depends_on:
- image-build-amd
@@ -153,7 +139,8 @@ steps:
- vllm/model_executor/models/qwen3_dflash.py
- tests/v1/spec_decode/test_speculators_correctness.py
commands:
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- label: Spec Decode MTP hybrid (B200)
timeout_in_minutes: 30
@@ -1,35 +0,0 @@
---
name: ci-fails-buildkite
description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
---
# Diagnosing vLLM Buildkite CI Failures
Buildkite logs are public; no login needed.
`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
## Fetching logs
```bash
# All failed jobs in a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# All failed jobs in a build (--soft also includes soft-failed jobs;
# --all fetches every finished job):
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
# work; pass "-" as a second argument to stream to stdout:
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
```
To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
```bash
./ci-clean-log.sh ci.log
```
## Reference
See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.
+15 -9
View File
@@ -2,14 +2,15 @@
# for more info about CODEOWNERS file
# This lists cover the "core" components of vLLM that require careful review
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
/vllm/lora @jeejeelee
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
/vllm/model_executor/layers/mamba @tdoublep @tomeras91
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
/vllm/model_executor/model_loader @22quinn
/vllm/model_executor/layers/batch_invariant.py @yewentao256
/vllm/ir @ProExpertProg
@@ -22,14 +23,9 @@
# Any change to the VllmConfig changes can have a large user-facing impact,
# so spam a lot of people
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
/vllm/config/cache.py @heheda12345
# Config utils
/vllm/config/utils.py @hmellor
/vllm/engine/arg_utils.py @hmellor
/vllm/utils/argparse_utils.py
# Entrypoints
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
@@ -79,7 +75,7 @@
/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche
# Model runner V2
/vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256
/vllm/v1/worker/gpu @WoosukKwon @njhill
/vllm/v1/worker/gpu/kv_connector.py @orozery
# CI & building
@@ -119,6 +115,16 @@
/vllm/model_executor/models/transformers @hmellor
/tests/models/test_transformers.py @hmellor
# Observability
/vllm/config/observability.py @markmc
/vllm/v1/metrics @markmc
/tests/v1/metrics @markmc
/vllm/tracing.py @markmc
/tests/v1/tracing/test_tracing.py @markmc
/vllm/config/kv_events.py @markmc
/vllm/distributed/kv_events.py @markmc
/tests/distributed/test_events.py @markmc
# Docs
/docs/mkdocs @hmellor
/docs/**/*.yml @hmellor
-5
View File
@@ -1,5 +0,0 @@
# Custom self-hosted runner labels (e.g. the autoscaling vllm-runners pool) so
# actionlint doesn't flag them as unknown in `runs-on`.
self-hosted-runner:
labels:
- vllm-runners
+1
View File
@@ -21,6 +21,7 @@ updates:
- dependency-name: "torchvision"
- dependency-name: "xformers"
- dependency-name: "lm-format-enforcer"
- dependency-name: "gguf"
- dependency-name: "compressed-tensors"
- dependency-name: "ray[cgraph]" # Ray Compiled Graph
- dependency-name: "lm-eval"
+4 -8
View File
@@ -144,12 +144,12 @@ pull_request_rules:
- label != stale
- or:
- files~=^examples/.*mistral.*\.py
- files~=^tests/.*(?:mistral|voxtral|mixtral|pixtral).*\.py
- files~=^vllm/model_executor/models/.*(?:mistral|voxtral|mixtral|pixtral).*\.py
- files~=^tests/.*mistral.*\.py
- files~=^vllm/model_executor/models/.*mistral.*\.py
- files~=^vllm/reasoning/.*mistral.*\.py
- files~=^vllm/tool_parsers/.*mistral.*\.py
- files~=^vllm/transformers_utils/.*(?:mistral|voxtral|pixtral).*\.py
- title~=(?i)(?:mistral|ministral|voxtral|mixtral|pixtral)
- files~=^vllm/transformers_utils/.*mistral.*\.py
- title~=(?i)Mistral
actions:
label:
add:
@@ -388,13 +388,9 @@ pull_request_rules:
- or:
- files~=^tests/tool_use/
- files~=^tests/tool_parsers/
- files~=^tests/parser/
- files~=^tests/reasoning/
- files~=^tests/entrypoints/openai/.*tool.*
- files~=^tests/entrypoints/anthropic/.*tool.*
- files~=^vllm/tool_parsers/
- files~=^vllm/parser/
- files~=^vllm/reasoning/
- files=docs/features/tool_calling.md
- files~=^examples/tool_calling/
actions:
+1 -5
View File
@@ -46,16 +46,12 @@ jobs:
pre-commit:
needs: pre-run-check
if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
runs-on: [self-hosted, linux, x64, vllm-runners]
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
with:
python-version: "3.12"
# Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its
# wget + tar -xJ self-download, which the self-hosted runner image lacks
# (no wget/xz). Pinned to shellcheck 0.10.0 to match the script's "stable".
- run: python -m pip install shellcheck-py==0.10.0.1
- run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
- run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json"
- run: echo "::add-matcher::.github/workflows/matchers/mypy.json"
+2 -5
View File
@@ -9,7 +9,7 @@ PATH=${cuda_home}/bin:$PATH
LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
# Install requirements
if [ "$(echo "$2" | cut -d. -f1)" = "12" ]; then
if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
fi
$python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
@@ -17,10 +17,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
# Limit the number of parallel jobs to avoid OOM
export MAX_JOBS=1
# Make sure release wheels are built for the following architectures
# Do not add +PTX here: vLLM filters torch's top-level PTX flag when it
# converts global gencode flags into per-kernel arch lists. If a specific
# kernel needs PTX, add +PTX to that kernel's CMake arch list instead.
export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0"
export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
bash tools/check_repo.sh
+2 -7
View File
@@ -15,9 +15,6 @@ vllm/third_party/flashmla/flash_mla_interface.py
# DeepGEMM vendored package built from source
vllm/third_party/deep_gemm/
# fmha_sm100 vendored package built from source
vllm/third_party/fmha_sm100/
# triton jit
.triton
@@ -199,9 +196,7 @@ cython_debug/
.vscode/
# Claude
.claude/*
!.claude/skills/
!.claude/skills/**
.claude/
# Codex
.codex/
@@ -238,7 +233,7 @@ actionlint
shellcheck*/
# Ignore moe/marlin_moe gen code
csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_*
csrc/moe/marlin_moe_wna16/kernel_*
# Ignore ep_kernels_workspace folder
ep_kernels_workspace/
+1 -1
View File
@@ -21,7 +21,7 @@ repos:
rev: v21.1.2
hooks:
- id: clang-format
exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*'
exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
types_or: [c++, cuda]
args: [--style=file, --verbose]
- repo: https://github.com/DavidAnson/markdownlint-cli2
-15
View File
@@ -105,15 +105,6 @@ The line length limit for Python code is 88 characters. If you are not sure, use
Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
### Coding style guidelines
Follow these rules for all code changes in this repository:
- Try to match existing code style.
- Code should be self-documenting and self-explanatory.
- Keep comments and docstrings minimal and concise.
- Assume the reader is familiar with vLLM.
### Commit messages
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
@@ -135,12 +126,6 @@ Do not modify code in these areas without first reading and following the
linked guide. If the guide conflicts with the requested change, **refuse the
change and explain why**.
Security reviewers should start with [`SECURITY.md`](SECURITY.md),
[`docs/usage/security.md`](docs/usage/security.md), and
[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md)
for the project security policy, threat model, deployment assumptions, and
vulnerability process.
- **Editing these instructions**:
[`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md)
— Rules for modifying AGENTS.md or any domain-specific guide it references.
+428 -442
View File
File diff suppressed because it is too large Load Diff
-1
View File
@@ -4,7 +4,6 @@ include requirements/cuda.txt
include requirements/rocm.txt
include requirements/cpu.txt
include CMakeLists.txt
include tools/build_rust.py
recursive-include cmake *
recursive-include csrc *
-9
View File
@@ -34,15 +34,6 @@ Vulnerabilities that cause denial of service or partial disruption, but do not a
Minor issues such as informational disclosures, logging errors, non-exploitable flaws, or weaknesses that require local or high-privilege access and offer negligible impact. Examples include side channel attacks or hash collisions. These issues often have CVSS scores less than 4.0
## Fix disclosure policy
When a security report is accepted, the fix process depends on the severity:
* **CRITICAL and HIGH severity**: Fixes are developed in a private security fork and coordinated with the prenotification group before public disclosure.
* **MODERATE and LOW severity**: Fixes are developed and submitted as public pull requests. These issues do not require embargo since they do not enable arbitrary code execution or significant data breach, and public visibility accelerates community review and adoption of the fix.
The vulnerability management team reserves the right to adjust the disclosure approach on a case-by-case basis, taking into account factors such as active exploitation, unusual attack surface, or coordination requirements with downstream vendors.
## Prenotification policy
For certain security issues of CRITICAL, HIGH, or MODERATE severity level, we may prenotify certain organizations or vendors that ship vLLM. The purpose of this prenotification is to allow for a coordinated release of fixes for severe issues.
+13 -9
View File
@@ -108,6 +108,7 @@ python benchmark.py \
--backends flash triton flashinfer \
--batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
--num-layers 10 \
--repeats 5 \
--output-csv results.csv
```
@@ -163,17 +164,14 @@ python benchmark.py \
# Model configuration
--num-layers N # Number of layers
--head-dim N # Head dimension
--v-head-dim N # Value head dimension (defaults to --head-dim)
--num-q-heads N # Query heads
--num-kv-heads N # KV heads
--block-size N # Block size
--kv-lora-rank N # MLA KV LoRA rank
--qk-nope-head-dim N # MLA non-RoPE QK head dim
--qk-rope-head-dim N # MLA RoPE QK head dim
# Benchmark settings
--device DEVICE # Device (default: cuda:0)
--warmup-ms N # Warmup window in ms for triton do_bench
--repeats N # Repetitions
--warmup-iters N # Warmup iterations
--profile-memory # Profile memory usage
# Parameter sweeps
@@ -213,6 +211,8 @@ config = BenchmarkConfig(
num_kv_heads=1,
block_size=128,
device="cuda:0",
repeats=5,
warmup_iters=3,
)
# CUTLASS MLA with specific num_kv_splits
@@ -253,10 +253,14 @@ formatter.save_json(results, "output.json")
## Tips
**1. Save results** - Always use `--output-csv` or `--output-json`
**1. Warmup matters** - Use `--warmup-iters 10` for stable results
**2. Test incrementally** - Start with `--num-layers 1`
**2. Multiple repeats** - Use `--repeats 20` for low variance
**3. Extended grammar** - Leverage spec decode, chunked prefill patterns
**3. Save results** - Always use `--output-csv` or `--output-json`
**4. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values
**4. Test incrementally** - Start with `--num-layers 1 --repeats 1`
**5. Extended grammar** - Leverage spec decode, chunked prefill patterns
**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values
+40 -261
View File
@@ -26,9 +26,6 @@ Examples:
"""
import argparse
import os
import shutil
import subprocess
import sys
from dataclasses import replace
from pathlib import Path
@@ -53,16 +50,6 @@ from common import (
from vllm.v1.worker.workspace import init_workspace_manager
def _str2bool(v) -> bool:
if isinstance(v, bool):
return v
if v.lower() in ("true", "1", "yes", "t"):
return True
if v.lower() in ("false", "0", "no", "f"):
return False
raise argparse.ArgumentTypeError(f"expected a boolean, got {v!r}")
def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
"""Run standard attention benchmark (Flash/Triton/FlashInfer)."""
from runner import run_attention_benchmark
@@ -96,15 +83,13 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
else:
return run_standard_attention_benchmark(config)
except Exception as e:
error_msg = str(e) or repr(e)
return BenchmarkResult(
config=config,
mean_time=float("inf"),
median_time=float("inf"),
std_time=0,
min_time=float("inf"),
max_time=float("inf"),
error=error_msg,
error=str(e),
)
@@ -130,12 +115,9 @@ def run_model_parameter_sweep(
"""
all_results = []
sweep_desc = (
f"{sweep.param_name} = {sweep.values}"
if sweep.param_name
else f"{len(sweep.values)} configurations"
console.print(
f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]"
)
console.print(f"[yellow]Model sweep mode: testing {sweep_desc}[/]")
total = len(backends) * len(batch_specs) * len(sweep.values)
@@ -143,9 +125,9 @@ def run_model_parameter_sweep(
for backend in backends:
for spec in batch_specs:
for value in sweep.values:
# Create config with modified model parameter(s)
# Create config with modified model parameter
config_args = base_config_args.copy()
sweep.apply(config_args, value)
config_args[sweep.param_name] = value
# Create config with original backend for running
clean_config = BenchmarkConfig(
@@ -162,21 +144,13 @@ def run_model_parameter_sweep(
all_results.append(result)
if not result.success:
err_label = (
f"{sweep.param_name}={value}"
if sweep.param_name
else f"{value}"
)
console.print(
f"[red]Error {backend} {spec} {err_label}"
f": {result.error}[/]"
f"[red]Error {backend} {spec} {sweep.param_name}="
f"{value}: {result.error}[/]"
)
pbar.update(1)
if base_config_args.get("ncu_profile"):
return all_results
# Display sweep results - create separate table for each parameter value
console.print("\n[bold green]Model Parameter Sweep Results:[/]")
formatter = ResultsFormatter(console)
@@ -210,10 +184,7 @@ def run_model_parameter_sweep(
)
for param_value in sorted_param_values:
label = (
f"{sweep.param_name} = {param_value}" if sweep.param_name else param_value
)
console.print(f"\n[bold cyan]{label}[/]")
console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]")
param_results = by_param_value[param_value]
# Create modified results with original backend names
@@ -229,9 +200,8 @@ def run_model_parameter_sweep(
formatter.print_table(modified_results, backends, compare_to_fastest=True)
# Show optimal backend for each (param_value, batch_spec) combination
sweep_name = sweep.param_name or "config"
console.print(
f"\n[bold cyan]Optimal backend for each ({sweep_name}, batch_spec):[/]"
f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]"
)
# Group by (param_value, batch_spec)
@@ -266,10 +236,7 @@ def run_model_parameter_sweep(
for param_value, spec in sorted_keys:
# Print header when param value changes
if param_value != current_param_value:
header = (
f"{sweep.param_name}={param_value}" if sweep.param_name else param_value
)
console.print(f"\n [bold]{header}:[/]")
console.print(f"\n [bold]{sweep.param_name}={param_value}:[/]")
current_param_value = param_value
results = by_param_and_spec[(param_value, spec)]
@@ -355,9 +322,6 @@ def run_parameter_sweep(
pbar.update(1)
if base_config_args.get("ncu_profile"):
return all_results
# Display sweep results
console.print("\n[bold green]Sweep Results:[/]")
backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]
@@ -495,20 +459,6 @@ def main():
help="Prefill backends to compare (fa2, fa3, fa4). "
"Uses the first decode backend for impl construction.",
)
parser.add_argument(
"--fp8-output-scale",
type=float,
help="Static per-tensor scale enabling the MLA prefill FP8-output "
"comparison on FA4 (fused write vs standalone post-quant).",
)
parser.add_argument(
"--fuse-quant-op",
nargs="+",
type=_str2bool,
help="FP8-output write path(s) to run: false = bf16 attention + "
"standalone static-FP8 quant, true = FA4 writes FP8 directly. "
"Default: both.",
)
# Batch specifications
parser.add_argument(
@@ -524,35 +474,11 @@ def main():
parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
parser.add_argument("--block-size", type=int, default=16, help="Block size")
parser.add_argument(
"--v-head-dim",
type=int,
default=None,
help="Value head dimension (defaults to --head-dim if unset)",
)
# MLA-specific model dimensions
parser.add_argument(
"--kv-lora-rank", type=int, default=None, help="MLA KV LoRA rank"
)
parser.add_argument(
"--qk-nope-head-dim", type=int, default=None, help="MLA non-RoPE QK head dim"
)
parser.add_argument(
"--qk-rope-head-dim", type=int, default=None, help="MLA RoPE QK head dim"
)
# Benchmark settings
parser.add_argument("--device", default="cuda:0", help="Device")
parser.add_argument(
"--warmup-ms",
type=int,
default=None,
help=(
"Warmup window in ms for triton's do_bench (default: triton's own). "
"Has no effect with CUDA graphs; pass --no-cuda-graphs to use it."
),
)
parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
parser.add_argument(
"--kv-cache-dtype",
@@ -565,33 +491,10 @@ def main():
action=argparse.BooleanOptionalAction,
default=True,
help=(
"Use triton do_bench_cudagraph (True) or do_bench (False) "
"for timing. CUDA graphs eliminate CPU launch overhead "
"(default: True)"
"Launch kernels with CUDA graphs to eliminate CPU overhead"
"in measurements (default: True)"
),
)
parser.add_argument(
"--num-splits",
type=int,
default=None,
help="FlashAttention split-K factor (0=auto heuristic, 1=disabled, >1=force N)",
)
parser.add_argument(
"--ncu-profile",
action="store_true",
default=False,
help=(
"Enable Nsight Compute profiling mode. Automatically wraps the "
"script with ncu, capturing a profile with source correlation. "
"Use --ncu-output to set the output file name."
),
)
parser.add_argument(
"--ncu-output",
type=str,
default="profile",
help="Output file name for ncu profile (default: 'profile').",
)
# Parameter sweep (use YAML config for advanced sweeps)
parser.add_argument(
@@ -642,12 +545,6 @@ def main():
# Prefill backends (e.g., ["fa3", "fa4"])
args.prefill_backends = yaml_config.get("prefill_backends", None)
# FP8 output benchmark knobs; CLI wins.
if args.fp8_output_scale is None:
args.fp8_output_scale = yaml_config.get("fp8_output_scale", None)
if args.fuse_quant_op is None:
args.fuse_quant_op = yaml_config.get("fuse_quant_op", None)
# Check for special modes
args.mode = yaml_config.get("mode", None)
@@ -679,28 +576,23 @@ def main():
model = yaml_config["model"]
args.num_layers = model.get("num_layers", args.num_layers)
args.head_dim = model.get("head_dim", args.head_dim)
args.v_head_dim = model.get("v_head_dim", args.v_head_dim)
args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
args.block_size = model.get("block_size", args.block_size)
# MLA-specific dimensions
args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
args.qk_rope_head_dim = model.get("qk_rope_head_dim", args.qk_rope_head_dim)
# Benchmark settings (top-level keys)
if "device" in yaml_config:
args.device = yaml_config["device"]
if "warmup_ms" in yaml_config:
args.warmup_ms = yaml_config["warmup_ms"]
if "repeats" in yaml_config:
args.repeats = yaml_config["repeats"]
if "warmup_iters" in yaml_config:
args.warmup_iters = yaml_config["warmup_iters"]
if "profile_memory" in yaml_config:
args.profile_memory = yaml_config["profile_memory"]
if "kv_cache_dtype" in yaml_config:
args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
if "cuda_graphs" in yaml_config:
args.cuda_graphs = yaml_config["cuda_graphs"]
if "ncu_profile" in yaml_config:
args.ncu_profile = yaml_config["ncu_profile"]
# Parameter sweep configuration
if "parameter_sweep" in yaml_config:
@@ -720,7 +612,7 @@ def main():
if "model_parameter_sweep" in yaml_config:
sweep_config = yaml_config["model_parameter_sweep"]
args.model_parameter_sweep = ModelParameterSweep(
param_name=sweep_config.get("param_name"),
param_name=sweep_config["param_name"],
values=sweep_config["values"],
label_format=sweep_config.get(
"label_format", "{backend}_{param_name}_{value}"
@@ -739,32 +631,6 @@ def main():
console.print()
# Re-exec under ncu if --ncu-profile and not already inside ncu. This runs
# after YAML processing so ncu_profile set via config file is honored.
if args.ncu_profile and "_NCU_INNER" not in os.environ:
ncu = shutil.which("ncu")
if ncu is None:
print("Error: 'ncu' not found in PATH", file=sys.stderr)
sys.exit(1)
cmd = [
ncu,
"--profile-from-start",
"off",
"--set",
"full",
"--import-source",
"yes",
"-o",
args.ncu_output,
sys.executable,
*sys.argv,
]
env = os.environ.copy()
env["CUTE_DSL_LINEINFO"] = "1"
env["_NCU_INNER"] = "1"
print(f"Launching: {' '.join(cmd)}")
sys.exit(subprocess.call(cmd, env=env))
# Handle CLI-based parameter sweep (if not from YAML)
if (
(not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)
@@ -789,18 +655,6 @@ def main():
console.print(f"Batch specs: {', '.join(args.batch_specs)}")
console.print(f"KV cache dtype: {args.kv_cache_dtype}")
console.print(f"CUDA graphs: {args.cuda_graphs}")
if args.warmup_ms is not None and args.cuda_graphs:
console.print(
"[yellow]Warning: --warmup-ms is ignored with CUDA graphs "
"(do_bench_cudagraph warms up internally). Pass --no-cuda-graphs "
"to use it.[/]"
)
if args.num_splits == 0 and args.cuda_graphs:
console.print(
"[yellow]Warning: --num-splits 0 (FA3 heuristic) is not CUDA-graph "
"compatible and may fail or fall back. Pass --no-cuda-graphs or use "
"--num-splits >=1.[/]"
)
console.print()
init_workspace_manager(args.device)
@@ -808,68 +662,8 @@ def main():
# Run benchmarks
all_results = []
# Under ncu profiling the kernels run only to be captured by the profiler;
# timings are placeholder zeros, so the result tables and saved metrics are
# skipped. The Nsight Compute report (--ncu-output) holds the real data.
if args.ncu_profile:
console.print(
"[dim]ncu profiling enabled: result tables and saved metrics are "
"skipped (timings are placeholder zeros).[/]"
)
# FA4 fused FP8 output vs standalone post-quant, on the same fa4 kernel:
# the delta is the post-quant kernel the fused path removes.
fp8_output_scale = getattr(args, "fp8_output_scale", None)
if fp8_output_scale is not None:
decode_backend = backends[0]
fuse_variants = args.fuse_quant_op or [False, True]
label_of = {False: "post_quant", True: "fused"}
console.print(
f"[yellow]FP8 output comparison @ scale={fp8_output_scale} "
f"(prefill=fa4, decode impl={decode_backend})[/]"
)
fp8_results = []
total = len(fuse_variants) * len(args.batch_specs)
with tqdm(total=total, desc="FP8 output benchmarking") as pbar:
for spec in args.batch_specs:
for fuse in fuse_variants:
config = BenchmarkConfig(
backend=decode_backend,
batch_spec=spec,
num_layers=args.num_layers,
head_dim=args.head_dim,
num_q_heads=args.num_q_heads,
num_kv_heads=args.num_kv_heads,
block_size=args.block_size,
device=args.device,
repeats=args.repeats,
warmup_iters=args.warmup_iters,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
use_cuda_graphs=args.cuda_graphs,
prefill_backend="fa4",
)
result = run_benchmark(
config, output_scale=fp8_output_scale, fuse_quant_op=fuse
)
label = label_of[fuse]
labeled_config = replace(result.config, backend=label)
result = replace(result, config=labeled_config)
fp8_results.append(result)
if not result.success:
console.print(f"[red]Error {label} {spec}: {result.error}[/]")
pbar.update(1)
console.print("\n[bold green]FP8 Output Results:[/]")
formatter = ResultsFormatter(console)
labels = [label_of[f] for f in fuse_variants]
formatter.print_table(fp8_results, labels, compare_to_fastest=True)
all_results = fp8_results
# Handle special mode: decode_vs_prefill comparison
elif hasattr(args, "mode") and args.mode == "decode_vs_prefill":
if hasattr(args, "mode") and args.mode == "decode_vs_prefill":
console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
console.print(
"[dim]For each query length, testing both decode and prefill pipelines[/]"
@@ -914,11 +708,11 @@ def main():
num_kv_heads=args.num_kv_heads,
block_size=args.block_size,
device=args.device,
repeats=args.repeats,
warmup_iters=args.warmup_iters,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
use_cuda_graphs=args.cuda_graphs,
ncu_profile=args.ncu_profile,
warmup_ms=args.warmup_ms,
)
# Add decode pipeline config
@@ -955,7 +749,6 @@ def main():
result = BenchmarkResult(
config=config,
mean_time=timing["mean"],
median_time=timing.get("median", timing["mean"]),
std_time=timing["std"],
min_time=timing["min"],
max_time=timing["max"],
@@ -977,7 +770,6 @@ def main():
result = BenchmarkResult(
config=config,
mean_time=float("inf"),
median_time=float("inf"),
std_time=0,
min_time=float("inf"),
max_time=float("inf"),
@@ -987,9 +779,6 @@ def main():
pbar.update(1)
if args.ncu_profile:
return
# Display decode vs prefill results
console.print("\n[bold green]Decode vs Prefill Results:[/]")
@@ -1069,20 +858,15 @@ def main():
base_config_args = {
"num_layers": args.num_layers,
"head_dim": args.head_dim,
"v_head_dim": args.v_head_dim,
"num_q_heads": args.num_q_heads,
"num_kv_heads": args.num_kv_heads,
"block_size": args.block_size,
"device": args.device,
"repeats": args.repeats,
"warmup_iters": args.warmup_iters,
"profile_memory": args.profile_memory,
"kv_cache_dtype": args.kv_cache_dtype,
"use_cuda_graphs": args.cuda_graphs,
"ncu_profile": args.ncu_profile,
"warmup_ms": args.warmup_ms,
"num_splits": args.num_splits,
"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,
}
all_results = run_model_parameter_sweep(
backends,
@@ -1098,17 +882,15 @@ def main():
base_config_args = {
"num_layers": args.num_layers,
"head_dim": args.head_dim,
"v_head_dim": args.v_head_dim,
"num_q_heads": args.num_q_heads,
"num_kv_heads": args.num_kv_heads,
"block_size": args.block_size,
"device": args.device,
"repeats": args.repeats,
"warmup_iters": args.warmup_iters,
"profile_memory": args.profile_memory,
"kv_cache_dtype": args.kv_cache_dtype,
"use_cuda_graphs": args.cuda_graphs,
"ncu_profile": args.ncu_profile,
"warmup_ms": args.warmup_ms,
"num_splits": args.num_splits,
}
all_results = run_parameter_sweep(
backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -1132,17 +914,15 @@ def main():
batch_spec=spec,
num_layers=args.num_layers,
head_dim=args.head_dim,
v_head_dim=getattr(args, "v_head_dim", None),
num_q_heads=args.num_q_heads,
num_kv_heads=args.num_kv_heads,
block_size=args.block_size,
device=args.device,
repeats=args.repeats,
warmup_iters=args.warmup_iters,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
use_cuda_graphs=args.cuda_graphs,
ncu_profile=args.ncu_profile,
warmup_ms=args.warmup_ms,
num_splits=args.num_splits,
)
result = run_benchmark(config)
@@ -1155,10 +935,9 @@ def main():
pbar.update(1)
if not args.ncu_profile:
console.print("\n[bold green]Results:[/]")
formatter = ResultsFormatter(console)
formatter.print_table(decode_results, backends)
console.print("\n[bold green]Results:[/]")
formatter = ResultsFormatter(console)
formatter.print_table(decode_results, backends)
# Run prefill backend comparison
if prefill_backends:
@@ -1183,8 +962,9 @@ def main():
num_kv_heads=args.num_kv_heads,
block_size=args.block_size,
device=args.device,
repeats=args.repeats,
warmup_iters=args.warmup_iters,
profile_memory=args.profile_memory,
warmup_ms=args.warmup_ms,
prefill_backend=pb,
)
@@ -1200,17 +980,16 @@ def main():
pbar.update(1)
if not args.ncu_profile:
console.print("\n[bold green]Prefill Backend Results:[/]")
formatter = ResultsFormatter(console)
formatter.print_table(
prefill_results, prefill_backends, compare_to_fastest=True
)
console.print("\n[bold green]Prefill Backend Results:[/]")
formatter = ResultsFormatter(console)
formatter.print_table(
prefill_results, prefill_backends, compare_to_fastest=True
)
all_results = decode_results + prefill_results
# Save results (skip ncu profiling runs: timings are placeholder zeros)
if all_results and not args.ncu_profile:
# Save results
if all_results:
formatter = ResultsFormatter(console)
if args.output_csv:
formatter.save_csv(all_results, args.output_csv)
+6 -54
View File
@@ -15,8 +15,6 @@ from batch_spec import get_batch_type, parse_batch_spec
from rich.console import Console
from rich.table import Table
from vllm.triton_utils import triton
def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
"""
@@ -36,30 +34,6 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
return (0, 0, 0)
def run_do_bench(
benchmark_fn,
use_cuda_graphs: bool,
warmup_ms: int | None = None,
) -> list[float]:
kwargs: dict[str, Any] = {"return_mode": "all"}
if use_cuda_graphs:
result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)
else:
if warmup_ms is not None:
kwargs["warmup"] = warmup_ms
result = triton.testing.do_bench(benchmark_fn, **kwargs)
return result
def run_ncu_profile(benchmark_fn) -> None:
benchmark_fn()
torch.accelerator.synchronize()
torch.cuda.cudart().cudaProfilerStart()
benchmark_fn()
torch.accelerator.synchronize()
torch.cuda.cudart().cudaProfilerStop()
# Mock classes for vLLM attention infrastructure
@@ -208,37 +182,18 @@ class ParameterSweep:
@dataclass
class ModelParameterSweep:
"""Configuration for sweeping model configuration parameter(s).
"""Configuration for sweeping a model configuration parameter."""
Supports two modes:
- Single param: param_name="head_dim", values=[128, 256, 512]
- Multi param: values=[{head_dim: 192, v_head_dim: 128}, {head_dim: 256}]
When values are dicts, each dict's keys are applied as config overrides.
"""
param_name: str | None = None
values: list[Any] | None = None
label_format: str = "{backend}_{param_name}_{value}"
param_name: str # Name of the model config parameter to sweep (e.g., "num_q_heads")
values: list[Any] # List of values to test
label_format: str = "{backend}_{param_name}_{value}" # Result label template
def get_label(self, backend: str, value: Any) -> str:
"""Generate a label for a specific parameter value."""
if isinstance(value, dict):
return self.label_format.format(
backend=backend, param_name=self.param_name, value=value, **value
)
return self.label_format.format(
backend=backend, param_name=self.param_name, value=value
)
def apply(self, config_args: dict, value: Any) -> None:
"""Apply a sweep value to config args."""
if isinstance(value, dict):
config_args.update(value)
elif self.param_name is not None:
config_args[self.param_name] = value
else:
raise ValueError("param_name must be set if sweep values are not dicts")
@dataclass
class BenchmarkConfig:
@@ -253,10 +208,10 @@ class BenchmarkConfig:
block_size: int
device: str
dtype: torch.dtype = torch.float16
repeats: int = 1
warmup_iters: int = 3
profile_memory: bool = False
use_cuda_graphs: bool = False
ncu_profile: bool = False
warmup_ms: int | None = None
# "auto" or "fp8"
kv_cache_dtype: str = "auto"
@@ -271,7 +226,6 @@ class BenchmarkConfig:
# Backend-specific tuning
num_kv_splits: int | None = None # CUTLASS MLA
reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA
num_splits: int | None = None # FlashAttention split-K (0=auto, 1=disabled)
@dataclass
@@ -280,7 +234,6 @@ class BenchmarkResult:
config: BenchmarkConfig
mean_time: float # seconds
median_time: float # seconds
std_time: float # seconds
min_time: float # seconds
max_time: float # seconds
@@ -299,7 +252,6 @@ class BenchmarkResult:
return {
"config": asdict(self.config),
"mean_time": self.mean_time,
"median_time": self.median_time,
"std_time": self.std_time,
"min_time": self.min_time,
"max_time": self.max_time,
@@ -56,6 +56,8 @@ backends:
- TOKENSPEED_MLA # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)
device: "cuda:0"
repeats: 100
warmup_iters: 10
profile_memory: true
# Backend-specific tuning
@@ -1,44 +0,0 @@
# MLA prefill FP8-output microbenchmark (FA4).
# Compares the fused FP8 write against bf16 attention + a standalone static-FP8
# quant; the delta is the post-quant kernel the fused path removes.
# DeepSeek-Coder-V2-Lite dims; FA4 needs SM100/110.
#
# Usage:
# python benchmark.py --config configs/mla_fa4_fp8_output.yaml
description: "MLA prefill FA4 fused-FP8 output vs post-quant"
model:
name: "deepseek-v2-lite"
num_layers: 27
num_q_heads: 16
num_kv_heads: 1
head_dim: 576
kv_lora_rank: 512
qk_nope_head_dim: 128
qk_rope_head_dim: 64
v_head_dim: 128
block_size: 128
# Pure prefill (q_len == kv_len) so every token goes through forward_mha.
batch_specs:
- "q512"
- "q1k"
- "q2k"
- "q4k"
- "q8k"
- "2q4k"
- "4q4k"
- "8q4k"
# Only used to construct the MLA impl; the pure-prefill specs skip decode.
decode_backends:
- CUTLASS_MLA
# Sweep the two FP8 write paths (prefill backend is fixed to fa4).
fp8_output_scale: 0.1
fuse_quant_op: [false, true]
device: "cuda:0"
repeats: 50
warmup_iters: 10
@@ -51,6 +51,8 @@ backends:
- FLASHMLA # Hopper only
device: "cuda:0"
repeats: 5
warmup_iters: 3
profile_memory: true
# Analyze chunked prefill workspace size impact
@@ -124,3 +124,5 @@ prefill_backends:
- tokenspeed
device: "cuda:0"
repeats: 20
warmup_iters: 5
@@ -53,4 +53,6 @@ backends:
- FLASHINFER_MLA_SPARSE
device: "cuda:0"
repeats: 100
warmup_iters: 10
profile_memory: true
@@ -57,4 +57,6 @@ backends:
- FLASHINFER_MLA_SPARSE
device: "cuda:0"
repeats: 10
warmup_iters: 3
profile_memory: true
@@ -63,6 +63,8 @@ model:
# Benchmark settings
device: "cuda:0"
repeats: 15 # More repeats for spec decode variance
warmup_iters: 5
profile_memory: false
# Output
@@ -49,6 +49,8 @@ backends:
# Benchmark settings
device: "cuda:0"
repeats: 10 # More repeats for statistical significance
warmup_iters: 5
profile_memory: false
# Test these threshold values for optimization
@@ -43,4 +43,6 @@ backends:
- FLASHINFER
device: "cuda:0"
repeats: 5
warmup_iters: 3
profile_memory: false
@@ -1,142 +0,0 @@
# Standard attention decode benchmark configuration
# Sweeps num_q_heads and num_kv_heads to isolate effects of:
# 1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)
# 2. Absolute head count (fixed 4:1 ratio, vary scale)
model:
num_layers: 32
num_q_heads: 32 # Base value, overridden by sweep
num_kv_heads: 8 # Base value, overridden by sweep
head_dim: 128
block_size: 16
# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and
# head_dim where it differs from the base (128). Head counts are per-GPU
# (i.e. after TP sharding).
#
# Group A — vary GQA ratio (fixed q=32, head_dim=128):
# 32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)
#
# Groups B-E — real model configs at various TP degrees:
# Model head_dim Full TP2 TP4 TP8
# Llama 3 8B 128 32:8 16:4 8:2 4:1
# Llama 3 70B 128 64:8 32:4 16:2 8:1
# GPT-OSS 120B 64 64:8 32:4 16:2 8:1
# Llama 3 405B 128 128:8 64:4 32:2 16:1
model_parameter_sweep:
values:
# --- head_dim=128 (Llama 3 family) ---
- { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1
- { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1
- { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8
- { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4
- { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2
- { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1
- { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8
- { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4
- { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1
- { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1
- { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8
- { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4
- { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2
- { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1
# --- head_dim=64 (GPT-OSS 120B) ---
- { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8
- { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4
- { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2
- { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1
label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"
batch_specs:
# ---- batch_size x seq_len grid (decode: q_len=1) ----
# Small grid for quick iteration. Uncomment for full sweep.
# Batch size 1
- "q1s1k"
- "q1s512"
- "q1s2k"
- "q1s4k"
- "q1s8k"
- "q1s16k"
- "q1s32k"
# Batch size 2
- "2q1s512"
- "2q1s1k"
- "2q1s2k"
- "2q1s4k"
- "2q1s8k"
- "2q1s16k"
- "2q1s32k"
# Batch size 4
- "4q1s512"
- "4q1s1k"
- "4q1s2k"
- "4q1s4k"
- "4q1s8k"
- "4q1s16k"
- "4q1s32k"
# Batch size 8
- "8q1s1k"
- "8q1s512"
- "8q1s2k"
- "8q1s4k"
- "8q1s8k"
- "8q1s16k"
- "8q1s32k"
# Batch size 16
- "16q1s512"
- "16q1s1k"
- "16q1s2k"
- "16q1s4k"
- "16q1s8k"
- "16q1s16k"
- "16q1s32k"
# Batch size 32
- "32q1s512"
- "32q1s1k"
- "32q1s2k"
- "32q1s4k"
- "32q1s8k"
- "32q1s16k"
- "32q1s32k"
# Batch size 64
- "64q1s1k"
- "64q1s512"
- "64q1s2k"
- "64q1s4k"
- "64q1s8k"
- "64q1s16k"
- "64q1s32k"
# Batch size 128
- "128q1s512"
- "128q1s1k"
- "128q1s2k"
- "128q1s4k"
- "128q1s8k"
- "128q1s16k"
- "128q1s32k"
# Batch size 256
- "256q1s1k"
- "256q1s512"
- "256q1s2k"
- "256q1s4k"
- "256q1s8k"
- "256q1s16k"
- "256q1s32k"
# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
backends:
- FLASH_ATTN
- TRITON_ATTN
- FLASHINFER
device: "cuda:0"
profile_memory: false
@@ -1,108 +0,0 @@
# Standard attention prefill benchmark configuration
# Sweeps num_q_heads and num_kv_heads to isolate effects of:
# 1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)
# 2. Absolute head count (fixed 4:1 ratio, vary scale)
model:
num_layers: 32
num_q_heads: 32 # Base value, overridden by sweep
num_kv_heads: 8 # Base value, overridden by sweep
head_dim: 128
block_size: 16
# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and
# head_dim where it differs from the base (128). Head counts are per-GPU
# (i.e. after TP sharding).
#
# Group A — vary GQA ratio (fixed q=32, head_dim=128):
# 32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)
#
# Groups B-E — real model configs at various TP degrees:
# Model head_dim Full TP2 TP4 TP8
# Llama 3 8B 128 32:8 16:4 8:2 4:1
# Llama 3 70B 128 64:8 32:4 16:2 8:1
# GPT-OSS 120B 64 64:8 32:4 16:2 8:1
# Llama 3 405B 128 128:8 64:4 32:2 16:1
model_parameter_sweep:
values:
# --- head_dim=128 (Llama 3 family) ---
- { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1
- { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1
- { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8
- { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4
- { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2
- { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1
- { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8
- { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4
- { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1
- { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1
- { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8
- { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4
- { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2
- { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1
# --- head_dim=64 (GPT-OSS 120B) ---
- { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8
- { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4
- { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2
- { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1
label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"
batch_specs:
# ---- batch_size x prefill_len grid (prefill: q_len == seq_len) ----
# Total tokens = batch_size * prefill_len, and prefill compute scales with
# prefill_len^2, so the largest cells are expensive. Trim batch sizes or
# lengths for quick iteration.
# Batch size 1
- "q512"
- "q1k"
- "q2k"
- "q4k"
- "q8k"
- "q16k"
- "q32k"
# Batch size 2
- "2q512"
- "2q1k"
- "2q2k"
- "2q4k"
- "2q8k"
- "2q16k"
- "2q32k"
# Batch size 4
- "4q512"
- "4q1k"
- "4q2k"
- "4q4k"
- "4q8k"
- "4q16k"
- "4q32k"
# Batch size 8
- "8q512"
- "8q1k"
- "8q2k"
- "8q4k"
- "8q8k"
- "8q16k"
- "8q32k"
# Batch size 16
- "16q512"
- "16q1k"
- "16q2k"
- "16q4k"
- "16q8k"
- "16q16k"
- "16q32k"
# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
backends:
- FLASH_ATTN
- TRITON_ATTN
- FLASHINFER
device: "cuda:0"
profile_memory: false
+42 -90
View File
@@ -8,8 +8,6 @@ This module provides helpers for running MLA backends without
needing full VllmConfig integration.
"""
import statistics
import numpy as np
import torch
from batch_spec import parse_batch_spec
@@ -19,8 +17,6 @@ from common import (
MockIndexer,
MockKVBProj,
MockLayer,
run_do_bench,
run_ncu_profile,
setup_mla_dims,
)
@@ -708,8 +704,6 @@ def _run_single_benchmark(
device: torch.device,
indexer=None,
kv_cache_dtype: str | None = None,
output_scale: float | None = None,
fuse_quant_op: bool = False,
) -> BenchmarkResult:
"""
Run a single benchmark iteration.
@@ -723,11 +717,6 @@ def _run_single_benchmark(
mla_dims: MLA dimension configuration
device: Target device
indexer: Optional MockIndexer for sparse backends
output_scale: Static per-tensor FP8 scale for prefill output. None
keeps the plain bf16 output (no quantization).
fuse_quant_op: With output_scale set, True lets the prefill kernel write
FP8 directly; False runs bf16 attention then a standalone static-FP8
quant. The delta isolates the saved post-quant kernel.
Returns:
BenchmarkResult with timing statistics
@@ -831,86 +820,63 @@ def _run_single_benchmark(
num_prefill, mla_dims, query_fmt, device, torch.bfloat16
)
# Prefill FP8 output: fused (kernel writes e4m3) vs separate post-quant.
prefill_fp8_output = None
prefill_output_scale = None
prefill_quant_op = None
if has_prefill and output_scale is not None:
from vllm.platforms import current_platform
prefill_output_scale = torch.tensor(
[output_scale], device=device, dtype=torch.float32
)
if fuse_quant_op:
prefill_fp8_output = torch.empty_like(
prefill_inputs["output"], dtype=current_platform.fp8_dtype()
)
else:
from vllm.model_executor.layers.quantization.input_quant_fp8 import (
QuantFP8,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
)
prefill_quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
fused_output = output_scale is not None and fuse_quant_op
# Build forward function (runs a single decode/prefill pass)
# Build forward function
def forward_fn():
results = []
if has_decode:
results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
if has_prefill:
out = impl.forward_mha(
prefill_inputs["q"],
prefill_inputs["k_c_normed"],
prefill_inputs["k_pe"],
kv_cache,
metadata,
prefill_inputs["k_scale"],
prefill_fp8_output if fused_output else prefill_inputs["output"],
prefill_output_scale if fused_output else None,
)
if fused_output:
out = prefill_fp8_output
elif prefill_quant_op is not None:
out, _ = prefill_quant_op(
prefill_inputs["output"], prefill_output_scale
results.append(
impl.forward_mha(
prefill_inputs["q"],
prefill_inputs["k_c_normed"],
prefill_inputs["k_pe"],
kv_cache,
metadata,
prefill_inputs["k_scale"],
prefill_inputs["output"],
)
results.append(out)
)
return results[0] if len(results) == 1 else tuple(results)
def benchmark_fn():
for _ in range(config.num_layers):
# Warmup
for _ in range(config.warmup_iters):
forward_fn()
torch.accelerator.synchronize()
# Optionally capture a CUDA graph after warmup.
# Graph replay eliminates CPU launch overhead so timings reflect pure
# kernel time.
if config.use_cuda_graphs:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
forward_fn()
benchmark_fn = graph.replay
else:
benchmark_fn = forward_fn
if config.ncu_profile:
run_ncu_profile(benchmark_fn)
return BenchmarkResult(
config=config,
mean_time=0.0,
median_time=0.0,
std_time=0.0,
min_time=0.0,
max_time=0.0,
throughput_tokens_per_sec=0.0,
)
# Benchmark
times = []
for _ in range(config.repeats):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
start.record()
for _ in range(config.num_layers):
benchmark_fn()
end.record()
# Convert ms to seconds per layer
times = [t / 1000.0 / config.num_layers for t in all_ms]
mean_time = statistics.mean(times)
torch.accelerator.synchronize()
elapsed_ms = start.elapsed_time(end)
times.append(elapsed_ms / 1000.0 / config.num_layers)
mean_time = float(np.mean(times))
return BenchmarkResult(
config=config,
mean_time=mean_time,
median_time=statistics.median(times),
std_time=statistics.stdev(times) if len(times) > 1 else 0.0,
min_time=min(times),
max_time=max(times),
std_time=float(np.std(times)),
min_time=float(np.min(times)),
max_time=float(np.max(times)),
throughput_tokens_per_sec=total_q / mean_time if mean_time > 0 else 0,
)
@@ -920,8 +886,6 @@ def _run_mla_benchmark_batched(
configs_with_params: list[tuple], # [(config, threshold, num_splits), ...]
index_topk: int = 2048,
prefill_backend: str | None = None,
output_scale: float | None = None,
fuse_quant_op: bool = False,
) -> list[BenchmarkResult]:
"""
Unified batched MLA benchmark runner for all backends.
@@ -1061,8 +1025,6 @@ def _run_mla_benchmark_batched(
device,
indexer=indexer,
kv_cache_dtype=kv_cache_dtype,
output_scale=output_scale,
fuse_quant_op=fuse_quant_op,
)
results.append(result)
@@ -1090,8 +1052,6 @@ def run_mla_benchmark(
num_kv_splits: int | None = None,
index_topk: int = 2048,
prefill_backend: str | None = None,
output_scale: float | None = None,
fuse_quant_op: bool = False,
) -> BenchmarkResult | list[BenchmarkResult]:
"""
Unified MLA benchmark runner for all backends.
@@ -1111,9 +1071,6 @@ def run_mla_benchmark(
index_topk: Topk value for sparse MLA backends (default 2048)
prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
When set, forces the specified FlashAttention version for prefill.
output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).
fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill
kernel vs a standalone post-quant kernel. See _run_single_benchmark.
Returns:
BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
@@ -1138,12 +1095,7 @@ def run_mla_benchmark(
# Use unified batched execution
results = _run_mla_benchmark_batched(
backend,
configs_with_params,
index_topk,
prefill_backend=prefill_backend,
output_scale=output_scale,
fuse_quant_op=fuse_quant_op,
backend, configs_with_params, index_topk, prefill_backend=prefill_backend
)
# Return single result or list based on input
+60 -53
View File
@@ -9,20 +9,13 @@ This module provides helpers for running standard attention backends
"""
import logging
import statistics
import types
from contextlib import contextmanager
import numpy as np
import torch
from batch_spec import parse_batch_spec, reorder_for_flashinfer
from common import (
BenchmarkConfig,
BenchmarkResult,
MockLayer,
get_attention_scale,
run_do_bench,
run_ncu_profile,
)
from common import BenchmarkConfig, BenchmarkResult, MockLayer, get_attention_scale
from vllm.config import (
CacheConfig,
@@ -215,13 +208,6 @@ def _create_backend_impl(
scale = get_attention_scale(config.head_dim)
# Set v_head_dim for diff-headdim backends. Always reset (defaulting to
# head_dim) so a prior run's value doesn't leak into this one via the
# backend's class-level state.
if hasattr(backend_class, "set_head_size_v"):
v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
backend_class.set_head_size_v(v_dim)
impl = backend_class.get_impl_cls()(
num_heads=config.num_q_heads,
head_size=config.head_dim,
@@ -314,7 +300,6 @@ def _create_input_tensors(
from vllm.platforms import current_platform
q_dtype = current_platform.fp8_dtype()
v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
q_list = [
torch.randn(
total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
@@ -328,7 +313,9 @@ def _create_input_tensors(
for _ in range(config.num_layers)
]
v_list = [
torch.randn(total_q, config.num_kv_heads, v_dim, device=device, dtype=dtype)
torch.randn(
total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
)
for _ in range(config.num_layers)
]
return q_list, k_list, v_list
@@ -402,17 +389,14 @@ def _run_single_benchmark(
device: torch.device,
dtype: torch.dtype,
) -> tuple:
"""Run single benchmark using triton's do_bench_cudagraph/do_bench.
Returns:
(timing_stats, mem_stats) where timing_stats is a dict with
mean/std/min/max in seconds per layer.
"""
"""Run single benchmark iteration with warmup and timing loop."""
total_q = q_list[0].shape[0]
v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
out = torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype)
out = torch.empty(
total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
)
def benchmark_fn():
# Warmup
for _ in range(config.warmup_iters):
for i in range(config.num_layers):
impl.forward(
layer,
@@ -423,22 +407,52 @@ def _run_single_benchmark(
attn_metadata,
output=out,
)
torch.accelerator.synchronize()
if config.ncu_profile:
run_ncu_profile(benchmark_fn)
timing_stats = dict.fromkeys(("mean", "median", "std", "min", "max"), 0.0)
# Optionally capture a CUDA graph after warmup.
# Graph replay eliminates CPU launch overhead so timings reflect pure
# kernel time.
if config.use_cuda_graphs:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
for i in range(config.num_layers):
impl.forward(
layer,
q_list[i],
k_list[i],
v_list[i],
cache_list[i],
attn_metadata,
output=out,
)
benchmark_fn = graph.replay
else:
all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
# Convert ms to seconds per layer
times = [t / 1000.0 / config.num_layers for t in all_ms]
timing_stats = {
"mean": statistics.mean(times),
"std": statistics.stdev(times) if len(times) > 1 else 0.0,
"min": min(times),
"max": max(times),
"median": statistics.median(times),
}
def benchmark_fn():
for i in range(config.num_layers):
impl.forward(
layer,
q_list[i],
k_list[i],
v_list[i],
cache_list[i],
attn_metadata,
output=out,
)
# Benchmark
times = []
for _ in range(config.repeats):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
benchmark_fn()
end.record()
torch.accelerator.synchronize()
elapsed_ms = start.elapsed_time(end)
times.append(elapsed_ms / 1000.0 / config.num_layers) # seconds per layer
mem_stats = {}
if config.profile_memory:
@@ -447,7 +461,7 @@ def _run_single_benchmark(
"reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
}
return timing_stats, mem_stats
return times, mem_stats
# ============================================================================
@@ -527,12 +541,6 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
common_attn_metadata=common_metadata,
)
# Override num_splits for split-K testing (FlashAttention only)
if config.num_splits is not None and hasattr(
attn_metadata, "max_num_splits"
):
attn_metadata.max_num_splits = config.num_splits
# Only quantize queries when the impl supports it
quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
impl, "supports_quant_query_input", False
@@ -545,7 +553,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
config, max_num_blocks, backend_class, device, dtype
)
timing_stats, mem_stats = _run_single_benchmark(
times, mem_stats = _run_single_benchmark(
config,
impl,
layer,
@@ -558,16 +566,15 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
dtype,
)
mean_time = timing_stats["mean"]
mean_time = np.mean(times)
throughput = total_q / mean_time if mean_time > 0 else 0
return BenchmarkResult(
config=config,
mean_time=mean_time,
median_time=timing_stats["median"],
std_time=timing_stats["std"],
min_time=timing_stats["min"],
max_time=timing_stats["max"],
std_time=np.std(times),
min_time=np.min(times),
max_time=np.max(times),
throughput_tokens_per_sec=throughput,
memory_allocated_mb=mem_stats.get("allocated_mb"),
memory_reserved_mb=mem_stats.get("reserved_mb"),
@@ -92,6 +92,7 @@ def run_baseline(
llm = LLM(
model=model,
enable_prefix_caching=False,
enable_chunked_prefill=False,
**extra_args,
)
sampling_params = SamplingParams(max_tokens=1)
@@ -193,6 +194,7 @@ async def _run_extraction_async(
engine_args = AsyncEngineArgs(
model=model,
enable_prefix_caching=False,
enable_chunked_prefill=False,
max_num_batched_tokens=40960,
max_model_len=40960,
speculative_config={
-358
View File
@@ -1,358 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Benchmark and regression-test pinned (page-locked) CPU memory for vLLM.
Verifies that enabling pinned memory does not regress throughput or latency
compared to unpinned memory. Each condition runs in an isolated ``spawn``
subprocess so both start from a cold CUDA context, giving an unbiased
comparison.
Usage
-----
Run all tests with the default model::
python benchmarks/benchmark_pin_memory.py -v
Override the model and optional max-model-len::
python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B -v
python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B \
--max-model-len 8192 -v
Run only throughput or latency tests::
python benchmarks/benchmark_pin_memory.py -v -k test_throughput
python benchmarks/benchmark_pin_memory.py -v -k test_latency
Run only the v1 or v2 runner variant::
python benchmarks/benchmark_pin_memory.py -v -k v1
python benchmarks/benchmark_pin_memory.py -v -k v2
Note: on WSL2, v1 runner tests are skipped because pin memory is not available
for the v1 runner without cpu_offload_gb. Run on other platforms to exercise v1.
"""
import argparse
import json
import multiprocessing
import sys
import tempfile
import pytest
# Allow up to 2% degradation. Both benchmark runs start from an identical
# cold CUDA context (separate spawn subprocesses), so the measured difference
# reflects the genuine pin_memory overhead rather than cold/warm ordering bias.
_THROUGHPUT_TOLERANCE = 0.98
_THROUGHPUT_NUM_REQUESTS = 200
_THROUGHPUT_INPUT_LEN = 128
_THROUGHPUT_OUTPUT_LEN = 512
_THROUGHPUT_MAX_NUM_SEQS = 128
# Latency benchmark constants — match latency.py defaults.
_LATENCY_TOLERANCE = 1.02 # Allow up to 2% latency regression.
_LATENCY_BATCH_SIZE = 64
_LATENCY_INPUT_LEN = 32
_LATENCY_OUTPUT_LEN = 128
_LATENCY_WARMUP_ITERS = 5
_LATENCY_BENCH_ITERS = 15
_DEFAULT_MODEL = "unsloth/Qwen3-1.7B"
_DEFAULT_MAX_MODEL_LEN = 16384
def _benchmark_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument("--model", default=_DEFAULT_MODEL)
parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
args, _ = parser.parse_known_args()
return args
@pytest.fixture
def model() -> str:
return _benchmark_args().model
@pytest.fixture
def max_model_len() -> int:
return _benchmark_args().max_model_len
def _skip_if_pin_memory_not_available(engine_args_kwargs: dict) -> None:
"""Skip the current pytest test if pin_memory is unavailable for this config."""
import vllm.utils.platform_utils as pu
from vllm.config import set_current_vllm_config
from vllm.engine.arg_utils import EngineArgs
vllm_config = EngineArgs(**engine_args_kwargs).create_engine_config()
with set_current_vllm_config(vllm_config):
pu.is_pin_memory_available.cache_clear()
if not pu.is_pin_memory_available():
import os
runner = "v2" if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") == "1" else "v1"
model = engine_args_kwargs.get("model", "unknown")
print(
f"\033[33mSKIP: pin_memory not available for "
f"{runner} runner, model={model}\033[0m"
)
pytest.skip("pin_memory not available for this configuration")
def _throughput_worker(
pin: bool,
engine_args_kwargs: dict,
q: "multiprocessing.Queue[float]",
v2_mode: bool = False,
) -> None:
"""Run throughput benchmark in a fresh spawn subprocess.
Delegates to vllm/benchmarks/throughput.py main() using the random dataset,
so the methodology matches the official benchmark. Results are written to a
temp JSON file and forwarded through the queue as tokens/s.
v2_mode: when True, monkeypatches is_uva_available() to always return True
so the v2 model runner's UVA buffers remain functional even when pin=False.
This isolates the non-UVA pin_memory paths in v2.
"""
import vllm.utils.platform_utils as pu
from vllm.platforms import current_platform
pu.is_pin_memory_available.cache_clear()
pu.is_uva_available.cache_clear()
type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
if v2_mode:
pu.is_uva_available = lambda: True
from vllm.benchmarks.throughput import add_cli_args
from vllm.benchmarks.throughput import main as throughput_main
parser = argparse.ArgumentParser()
add_cli_args(parser)
args = parser.parse_args([])
for key, val in engine_args_kwargs.items():
setattr(args, key, val)
args.max_num_seqs = _THROUGHPUT_MAX_NUM_SEQS
args.dataset_name = "random"
args.input_len = _THROUGHPUT_INPUT_LEN
args.output_len = _THROUGHPUT_OUTPUT_LEN
# Nullify defaults that conflict with explicit input/output_len.
args.random_input_len = None
args.random_output_len = None
args.random_prefix_len = None
args.num_prompts = _THROUGHPUT_NUM_REQUESTS
args.seed = 0
args.disable_detokenize = True
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
tmp_path = f.name
args.output_json = tmp_path
throughput_main(args)
with open(tmp_path) as f:
results = json.load(f)
q.put(results["tokens_per_second"])
def _run_throughput_benchmark(
pin: bool,
engine_args_kwargs: dict,
v2_mode: bool = False,
) -> float:
ctx = multiprocessing.get_context("spawn")
q = ctx.Queue()
p = ctx.Process(
target=_throughput_worker,
args=(pin, engine_args_kwargs, q, v2_mode),
)
p.start()
p.join()
if p.exitcode != 0:
raise RuntimeError(
f"Throughput benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
)
return q.get()
def _latency_worker(
pin: bool,
engine_args_kwargs: dict,
q: "multiprocessing.Queue[dict]",
v2_mode: bool = False,
) -> None:
"""Run latency benchmark in a fresh spawn subprocess.
Follows latency.py methodology: fixed batch of dummy token IDs, warmup
iterations to reach steady state, then timed iterations reduced to avg
and percentiles. Results are written to a temp JSON file by latency_main
and forwarded through the queue.
"""
import vllm.utils.platform_utils as pu
from vllm.platforms import current_platform
pu.is_pin_memory_available.cache_clear()
pu.is_uva_available.cache_clear()
type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
if v2_mode:
pu.is_uva_available = lambda: True
from vllm.benchmarks.latency import add_cli_args
from vllm.benchmarks.latency import main as latency_main
parser = argparse.ArgumentParser()
add_cli_args(parser)
args = parser.parse_args([])
for key, val in engine_args_kwargs.items():
setattr(args, key, val)
args.input_len = _LATENCY_INPUT_LEN
args.output_len = _LATENCY_OUTPUT_LEN
args.batch_size = _LATENCY_BATCH_SIZE
args.num_iters_warmup = _LATENCY_WARMUP_ITERS
args.num_iters = _LATENCY_BENCH_ITERS
args.profile = False
args.disable_detokenize = True
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
tmp_path = f.name
args.output_json = tmp_path
latency_main(args)
with open(tmp_path) as f:
results = json.load(f)
q.put(results)
def _run_latency_benchmark(
pin: bool,
engine_args_kwargs: dict,
v2_mode: bool = False,
) -> dict:
ctx = multiprocessing.get_context("spawn")
q = ctx.Queue()
p = ctx.Process(
target=_latency_worker,
args=(pin, engine_args_kwargs, q, v2_mode),
)
p.start()
p.join()
if p.exitcode != 0:
raise RuntimeError(
f"Latency benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
)
return q.get()
@pytest.mark.parametrize(
"test_v2_runner",
[
pytest.param(False, id="v1"),
pytest.param(True, id="v2"),
],
)
class TestPinnedMemory:
"""Verify pinned memory yields >= throughput vs unpinned via real vLLM inference."""
def test_throughput(self, monkeypatch, test_v2_runner, model, max_model_len):
"""Benchmark throughput with pin_memory forced on then off.
Delegates to vllm/benchmarks/throughput.py main() with the random
dataset. Each condition runs in an isolated spawn subprocess so both
start from a cold CUDA context, giving an unbiased comparison.
"""
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
engine_args_kwargs = dict(
model=model,
gpu_memory_utilization=0.88,
max_model_len=max_model_len,
enable_prefix_caching=False,
)
_skip_if_pin_memory_not_available(engine_args_kwargs)
unpinned_tps = _run_throughput_benchmark(
False, engine_args_kwargs, v2_mode=test_v2_runner
)
pinned_tps = _run_throughput_benchmark(
True, engine_args_kwargs, v2_mode=test_v2_runner
)
pct_diff = (pinned_tps - unpinned_tps) / unpinned_tps * 100
runner = "v2" if test_v2_runner else "v1"
print(
f"\n=== Throughput results ({runner} runner, {model}) ==="
f"\npin_memory=True: {pinned_tps:.1f} tok/s"
f"\npin_memory=False: {unpinned_tps:.1f} tok/s"
f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
)
assert pinned_tps >= unpinned_tps * _THROUGHPUT_TOLERANCE, (
f"Pinned throughput ({pinned_tps:.1f} tok/s) fell more than "
f"{(1.0 - _THROUGHPUT_TOLERANCE) * 100:.1f}% below "
f"unpinned ({unpinned_tps:.1f} tok/s)."
)
def test_latency(self, monkeypatch, test_v2_runner, model, max_model_len):
"""Benchmark per-batch latency with pin_memory forced on then off.
Follows vllm/benchmarks/latency.py: fixed dummy-token batch, warmup
iterations to reach steady state, then timed iterations reduced to avg
and percentiles. Subprocesses run serially so each gets a cold CUDA
context without GPU memory pressure from the other run.
"""
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
engine_args_kwargs = dict(
model=model,
gpu_memory_utilization=0.88,
max_model_len=max_model_len,
enable_prefix_caching=False,
)
_skip_if_pin_memory_not_available(engine_args_kwargs)
unpinned = _run_latency_benchmark(
False, engine_args_kwargs, v2_mode=test_v2_runner
)
pinned = _run_latency_benchmark(
True, engine_args_kwargs, v2_mode=test_v2_runner
)
pct_diff = (
(pinned["avg_latency"] - unpinned["avg_latency"])
/ unpinned["avg_latency"]
* 100
)
runner = "v2" if test_v2_runner else "v1"
print(
f"\n=== Latency results ({runner} runner, {model}) ==="
f"\npin_memory=True: avg={pinned['avg_latency']:.3f}s"
f" p50={pinned['percentiles']['50']:.3f}s"
f" p99={pinned['percentiles']['99']:.3f}s"
f"\npin_memory=False: avg={unpinned['avg_latency']:.3f}s"
f" p50={unpinned['percentiles']['50']:.3f}s"
f" p99={unpinned['percentiles']['99']:.3f}s"
f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
)
assert pinned["avg_latency"] <= unpinned["avg_latency"] * _LATENCY_TOLERANCE, (
f"Pinned avg latency ({pinned['avg_latency']:.3f}s) exceeded "
f"unpinned ({unpinned['avg_latency']:.3f}s) by more than "
f"{(_LATENCY_TOLERANCE - 1.0) * 100:.1f}%."
)
if __name__ == "__main__":
_parser = argparse.ArgumentParser(add_help=False)
_parser.add_argument("--model", default=_DEFAULT_MODEL)
_parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
_, _remaining = _parser.parse_known_args()
sys.exit(pytest.main([__file__] + _remaining))
@@ -1,277 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Copyright (c) 2025 FlyDSL Project Contributors
import json
import os
import torch
from aiter.test_common import run_perftest
from vllm.model_executor.layers.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
int4_w4a16_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
compressed_tensors_moe_w4a16_flydsl,
)
from vllm.platforms import current_platform
RoutingBuffers = tuple[
torch.Tensor, # sorted_token_ids
torch.Tensor, # sorted_weights
torch.Tensor, # sorted_expert_ids
torch.Tensor, # num_valid_ids (shape [1], i32)
int, # sorted_size
int, # blocks
]
MODEL_PARAMS_TO_TUNE = [
# (num_experts, inter_dim, hidden_size, topk)
(384, 256, 7168, 8), # Kimi K2.5 TP=8
(384, 512, 7168, 8), # Kimi K2.5 TP=4
]
NUM_TOKENS_TO_TUNE = [
1,
2,
4,
8,
16,
24,
32,
48,
64,
128,
256,
512,
1024,
2048,
4096,
8192,
]
TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256]
TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256]
TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256]
TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
TILE_CONFIGS = []
for tile_m in TILE_M_SEARCH_SPACE:
for tile_n in TILE_N_SEARCH_SPACE:
for tile_k in TILE_K_SEARCH_SPACE:
for tile_n2 in TILE_N2_SEARCH_SPACE:
for tile_k2 in TILE_K2_SEARCH_SPACE:
TILE_CONFIGS.append(
{
"tile_m": tile_m,
"tile_n": tile_n,
"tile_k": tile_k,
"tile_n2": tile_n2,
"tile_k2": tile_k2,
}
)
def tune_flydsl_moe_w4a16(
device: str = "cuda", num_iters: int = 100, num_warmup: int = 10
):
packed_factor = 8
w13_num_shards = 2
params_dtype = torch.bfloat16
group_size = 32
scale_factor = 0.01
for model_params in MODEL_PARAMS_TO_TUNE:
num_experts = model_params[0]
inter_dim = model_params[1]
hidden_size = model_params[2]
topk = model_params[3]
print(
f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, "
f"hidden_size={hidden_size}, topk={topk}...\n"
)
w2_scales_size = inter_dim
num_groups_w2 = w2_scales_size // group_size
num_groups_w13 = hidden_size // group_size
w13_weight = torch.randint(
0,
255,
(num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim),
dtype=torch.int32,
device=device,
)
w2_weight = torch.randint(
0,
255,
(num_experts, inter_dim // packed_factor, hidden_size),
dtype=torch.int32,
device=device,
)
w13_scale = scale_factor * torch.randn(
num_experts,
num_groups_w13,
w13_num_shards * inter_dim,
dtype=params_dtype,
device=device,
)
w2_scale = scale_factor * torch.randn(
num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device
)
w13 = w13_weight
w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13)
w13 = w13.view(-1).contiguous()
w2 = w2_weight
w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2)
w2 = w2.view(-1).contiguous()
w13_scale_flydsl = w13_scale
w2_scale_flydsl = w2_scale
if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
E, G, N = w13_scale.shape
w13_scale_flydsl = (
w13_scale_flydsl.view(E, G // 2, 2, N)
.permute(0, 1, 3, 2)
.contiguous()
.view(-1)
.contiguous()
)
elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
w13_scale_flydsl = w13_scale_flydsl.squeeze(1)
if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
E, G, N = w2_scale.shape
w2_scale_flydsl = (
w2_scale_flydsl.view(E, G // 2, 2, N)
.permute(0, 1, 3, 2)
.contiguous()
.view(-1)
.contiguous()
)
elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
w2_scale_flydsl = w2_scale_flydsl.squeeze(1)
w13_scale_flydsl = w13_scale_flydsl.contiguous()
w2_scale_flydsl = w2_scale_flydsl.contiguous()
w13.is_shuffled = True
w2.is_shuffled = True
w13_weight_scale = w13_scale.transpose(1, 2).contiguous()
w2_weight_scale = w2_scale.transpose(1, 2).contiguous()
w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8)
w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8)
moe_quant_config = int4_w4a16_moe_quant_config(
w1_scale=w13_weight_scale,
w2_scale=w2_weight_scale,
w1_zp=None,
w2_zp=None,
block_shape=[0, group_size],
)
tuned_config = {}
for num_tokens in NUM_TOKENS_TO_TUNE:
score = torch.rand(
(num_tokens, num_experts), device=device, dtype=torch.float32
)
topk_vals, topk_ids = torch.topk(score, k=topk, dim=1)
topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32)
x = torch.randn(
(num_tokens, hidden_size), dtype=torch.bfloat16, device=device
)
us_best = float("inf")
for tile_config in TILE_CONFIGS:
try:
tile_m = tile_config["tile_m"]
tile_n = tile_config["tile_n"]
tile_k = tile_config["tile_k"]
tile_n2 = tile_config["tile_n2"]
tile_k2 = tile_config["tile_k2"]
model_dim = x.shape[1]
assert model_dim % 64 == 0
assert model_dim % tile_k == 0
assert inter_dim % tile_n == 0
assert model_dim % tile_n2 == 0
assert inter_dim % tile_k2 == 0
assert ((tile_m * tile_k2) % 256) == 0
bytes_per_thread_x = (tile_m * tile_k2) // 256
assert (bytes_per_thread_x % 4) == 0
out, _us = run_perftest(
fused_flydsl_moe,
x,
w13,
w2,
num_experts,
inter_dim,
topk_weights,
topk_ids,
num_iters=num_iters,
num_warmup=num_warmup,
w1_scale=w13_scale_flydsl,
w2_scale=w2_scale_flydsl,
topk=topk_weights.shape[-1],
group_size=group_size,
doweight_stage1=False,
scale_is_bf16=True,
config=tile_config,
)
torch.accelerator.synchronize()
except Exception:
torch.accelerator.synchronize()
continue
else:
us = _us.item()
if us < us_best:
out_ref = fused_experts(
x,
w13_weight_packed,
w2_weight_packed,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SILU,
apply_router_weight_on_input=False,
global_num_experts=num_experts,
expert_map=None,
quant_config=moe_quant_config,
)
try:
assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1)
except Exception:
continue
else:
print(
f"For [num_tokens={num_tokens}, num_experts={num_experts}, " # noqa: E501
f"inter_dim={inter_dim}] found new best " # noqa: E501
f"config={tile_config}, us={us:0.3f}"
)
us_best = us
tuned_config[str(num_tokens)] = tile_config
device_name = current_platform.get_device_name().replace(" ", "_")
tuned_config_file_name = (
f"E={num_experts},N={inter_dim},device_name={device_name},"
f"dtype=int4_w4a16,backend=flydsl.json"
)
tuner_dir_path = os.path.dirname(os.path.realpath(__file__))
store_path = os.path.join(tuner_dir_path, tuned_config_file_name)
with open(store_path, "w") as f:
json.dump(tuned_config, f, indent=4)
print(
f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n" # noqa: E501
)
if __name__ == "__main__":
tune_flydsl_moe_w4a16(device="cuda")
-6
View File
@@ -792,12 +792,6 @@ def get_model_params(config):
topk = text_config.num_experts_per_tok
intermediate_size = text_config.moe_intermediate_size
hidden_size = text_config.hidden_size
elif architecture == "DiffusionGemmaForBlockDiffusion":
text_config = config.get_text_config()
E = text_config.num_experts
topk = text_config.top_k_experts
intermediate_size = text_config.moe_intermediate_size
hidden_size = text_config.hidden_size
elif architecture == "HunYuanMoEV1ForCausalLM":
E = config.num_experts
topk = config.moe_topk[0]
-248
View File
@@ -1,248 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Reproducible demonstration of the KV cache watermark (`--watermark`) for
# reducing preemption thrashing.
#
# The watermark is the fraction of total KV cache blocks the scheduler keeps
# free when admitting a waiting/preempted request into the running queue.
#
# Why this workload triggers thrashing:
# Requests are admitted based on the KV cache they need *at admission time*.
# With `--scheduler-reserve-full-isl` (default) the input length is reserved up
# front, but the *output* length is unknown and unreserved. A decode-heavy
# workload (output >> input) at high concurrency therefore over-admits while
# requests are short, then runs out of KV cache as they all grow during decode
# -> the scheduler preempts (recompute) recently-admitted requests, re-prefills
# them later, and repeats. The watermark keeps a block of KV cache free so
# running requests can grow into it instead of triggering this churn.
#
# This script launches `vllm serve` under a deliberately KV-constrained config
# and a decode-heavy workload, sweeping the watermark across several values, and
# reports the preemption count (scraped from /metrics), throughput, and latency
# percentiles for each. It then plots the results.
#
# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens
# (+/- 20% variance), sized to run each config for ~5 minutes.
#
# Usage:
# benchmarks/kv_cache_watermark.sh
# MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh
#
# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it).
set -euo pipefail
# ---- Config (override via environment) -------------------------------------
MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct}
TP=${TP:-1}
PORT=${PORT:-8000}
URL="http://127.0.0.1:${PORT}"
# Constrain the KV cache to a *near-critical* size: large enough that the engine
# can run stably, but small enough that greedy over-admission tips it into
# preemption thrashing. (Independent of GPU size, so the demo is reproducible.)
# At the default workload this fits ~1.5x the mean concurrent KV demand.
KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16}
MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192}
MAX_NUM_SEQS=${MAX_NUM_SEQS:-256}
# Optional weight loader (e.g. fastsafetensors on the GCP cluster).
LOAD_FORMAT=${LOAD_FORMAT:-auto}
# Decode-heavy workload: moderate input, long output, with length variance. The
# long output means preempted requests have generated a lot before eviction, so
# resuming them re-prefills a long sequence (high recomputation cost).
INPUT_LEN=${INPUT_LEN:-1000}
OUTPUT_LEN=${OUTPUT_LEN:-5000}
RANGE_RATIO=${RANGE_RATIO:-0.2}
CONCURRENCY=${CONCURRENCY:-128}
# Enough prompts to keep each config saturated for ~5+ minutes.
NUM_PROMPTS=${NUM_PROMPTS:-450}
OUTDIR=${OUTDIR:-./watermark_bench_results}
# Watermark fractions compared. "label value" per line; value=0 disables it.
CONFIGS=${CONFIGS:-"off 0
w0.02 0.02
w0.05 0.05
w0.10 0.10
w0.15 0.15"}
KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024))
mkdir -p "$OUTDIR"
SERVER_PID=""
cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; }
trap cleanup EXIT
scrape_preemptions() {
# Sum the vllm:num_preemptions_total counter across engines.
python - "${URL}/metrics" <<'PY'
import sys, urllib.request
total = 0.0
try:
body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace")
for line in body.splitlines():
if line.startswith("vllm:num_preemptions_total"):
total += float(line.rsplit(" ", 1)[-1])
except Exception as e: # noqa: BLE001
print(f"scrape error: {e}", file=sys.stderr)
print(int(total))
PY
}
wait_for_server() {
for _ in $(seq 1 300); do
if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
echo "ERROR: server process exited during startup" >&2; return 1
fi
sleep 5
done
echo "ERROR: server did not become ready" >&2; return 1
}
run_one() {
local label=$1 watermark=$2
echo
echo "==================== watermark: ${label} (${watermark}) ===================="
vllm serve "$MODEL" \
--tensor-parallel-size "$TP" \
--load-format "$LOAD_FORMAT" \
--kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \
--max-model-len "$MAX_MODEL_LEN" \
--max-num-seqs "$MAX_NUM_SEQS" \
--no-enable-prefix-caching \
--watermark "$watermark" \
--port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 &
SERVER_PID=$!
wait_for_server
sleep 5
local pre post
pre=$(scrape_preemptions)
vllm bench serve \
--backend vllm \
--base-url "$URL" \
--model "$MODEL" \
--dataset-name random \
--random-input-len "$INPUT_LEN" \
--random-output-len "$OUTPUT_LEN" \
--random-range-ratio "$RANGE_RATIO" \
--ignore-eos \
--num-prompts "$NUM_PROMPTS" \
--max-concurrency "$CONCURRENCY" \
--percentile-metrics "ttft,tpot,itl,e2el" \
--metric-percentiles "50,90,99" \
--save-result \
--result-dir "$OUTDIR" \
--result-filename "bench_${label}.json"
post=$(scrape_preemptions)
echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt"
kill "$SERVER_PID" 2>/dev/null || true
for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done
SERVER_PID=""
sleep 10
}
: >"${OUTDIR}/preemptions.txt"
while read -r label watermark; do
[[ -z "${label:-}" ]] && continue
run_one "$label" "$watermark"
done <<<"$CONFIGS"
echo
echo "==================== summary ===================="
python - "$OUTDIR" <<'PY'
import json, os, sys
outdir = sys.argv[1]
pre = {}
order = []
for line in open(os.path.join(outdir, "preemptions.txt")):
label, watermark, n = line.split()
pre[label] = (float(watermark), int(n))
order.append(label)
def g(d, *names):
for n in names:
if d.get(n) is not None:
return d[n]
return float("nan")
cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s",
"TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"]
print(" ".join(f"{c:>10}" for c in cols))
rows = []
for label in order:
watermark, n = pre[label]
d = json.load(open(os.path.join(outdir, f"bench_{label}.json")))
rows.append(dict(
label=label, watermark=watermark, preempt=n,
out_tok_s=g(d, "output_throughput"),
req_s=g(d, "request_throughput"),
ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"),
ttft_p99=g(d, "p99_ttft_ms"),
itl_p99=g(d, "p99_itl_ms"),
e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"),
))
print(" ".join(f"{str(v):>10}" for v in [
label, watermark, n,
f"{rows[-1]['out_tok_s']:.0f}",
f"{rows[-1]['req_s']:.3f}",
f"{rows[-1]['ttft_p50']/1000:.2f}",
f"{rows[-1]['ttft_p99']/1000:.2f}",
f"{rows[-1]['itl_p99']:.2f}",
f"{rows[-1]['e2el_p50']/1000:.1f}",
]))
print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)")
# ---- Plot -------------------------------------------------------------------
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
except Exception as e: # noqa: BLE001
print(f"\n(skip plot: matplotlib unavailable: {e})")
sys.exit(0)
x = [r["watermark"] for r in rows]
xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows]
idx = list(range(len(rows)))
fig, axes = plt.subplots(2, 2, figsize=(12, 8))
fig.suptitle(
f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}",
fontsize=12,
)
ax = axes[0][0]
ax.bar(idx, [r["preempt"] for r in rows], color="tab:red")
ax.set_title("Preemptions (lower is better)")
ax.set_ylabel("preemptions")
ax.set_xticks(idx); ax.set_xticklabels(xt)
ax = axes[0][1]
ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green")
ax.set_title("Output throughput (higher is better)")
ax.set_ylabel("tokens/s")
ax.set_xticks(idx); ax.set_xticklabels(xt)
ax = axes[1][0]
ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue")
ax.set_title("Inter-token latency p99 (lower is better)")
ax.set_ylabel("ITL p99 (ms)")
ax.set_xlabel("watermark fraction")
ax.set_xticks(idx); ax.set_xticklabels(xt)
ax = axes[1][1]
ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50")
ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99")
ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50")
ax.set_title("Latency (lower is better)")
ax.set_ylabel("seconds")
ax.set_xlabel("watermark fraction")
ax.set_xticks(idx); ax.set_xticklabels(xt)
ax.legend()
fig.tight_layout(rect=(0, 0, 1, 0.95))
out_png = os.path.join(outdir, "watermark_results.png")
fig.savefig(out_png, dpi=120)
print(f"\nWrote plot: {out_png}")
PY
@@ -65,32 +65,6 @@ class RequestArgs(NamedTuple):
limit_min_tokens: int # Use negative value for no limit
limit_max_tokens: int # Use negative value for no limit
timeout_sec: int
send_conversation_id: bool
headers: dict[str, str]
def parse_custom_header(header: str) -> tuple[str, str]:
separators = (":", "=")
for separator in separators:
if separator in header:
key, value = header.split(separator, 1)
key = key.strip()
value = value.strip()
if key:
return key, value
break
raise argparse.ArgumentTypeError(
"Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
)
def build_request_headers(
api_key: str | None, custom_headers: list[tuple[str, str]] | None
) -> dict[str, str]:
headers = dict(custom_headers or [])
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
class BenchmarkArgs(NamedTuple):
@@ -244,11 +218,12 @@ async def send_request(
max_tokens: int | None = None,
timeout_sec: int = 120,
conversation_id: str | None = None,
headers: dict[str, str] | None = None,
) -> ServerResponse:
payload = {
"model": model,
"messages": messages,
"seed": 0,
"temperature": 0.0,
}
if conversation_id is not None:
@@ -258,17 +233,13 @@ async def send_request(
payload["stream"] = True
payload["stream_options"] = {"include_usage": False}
# if min_tokens is not None:
# payload["min_tokens"] = min_tokens
if min_tokens is not None:
payload["min_tokens"] = min_tokens
if max_tokens is not None:
payload["max_tokens"] = max_tokens
request_headers = {"Content-Type": "application/json"}
if conversation_id is not None:
request_headers["X-Session-ID"] = str(conversation_id)
if headers is not None:
request_headers.update(headers)
headers = {"Content-Type": "application/json"}
# Calculate the timeout for the request
if max_tokens is not None:
@@ -294,7 +265,7 @@ async def send_request(
most_recent_timestamp: int = start_time
async with session.post(
url=chat_url, json=payload, headers=request_headers, timeout=timeout
url=chat_url, json=payload, headers=headers, timeout=timeout
) as response:
http_status = HTTPStatus(response.status)
if http_status == HTTPStatus.OK:
@@ -346,8 +317,6 @@ async def send_request(
latency = time.perf_counter_ns() - start_time
if ttft is None:
if stream:
valid_response = False
# The response was a single chunk
ttft = latency
@@ -454,8 +423,7 @@ async def send_turn(
min_tokens,
max_tokens,
req_args.timeout_sec,
conversation_id=conv_id if req_args.send_conversation_id else None,
headers=req_args.headers,
conversation_id=conv_id,
)
if response.valid is False:
@@ -904,7 +872,6 @@ def get_client_config(
# Arguments for API requests
chat_url = f"{args.url}/v1/chat/completions"
model_name = args.served_model_name if args.served_model_name else args.model
headers = build_request_headers(args.api_key, args.header)
req_args = RequestArgs(
chat_url=chat_url,
@@ -913,8 +880,6 @@ def get_client_config(
limit_min_tokens=args.limit_min_tokens,
limit_max_tokens=args.limit_max_tokens,
timeout_sec=args.request_timeout_sec,
send_conversation_id=args.send_conversation_id,
headers=headers,
)
return client_args, req_args
@@ -1280,19 +1245,19 @@ def process_statistics(
)
async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
async def get_server_info(url: str) -> None:
logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
async with aiohttp.ClientSession() as session:
# Get server version (not mandatory, "version" endpoint may not exist)
url_version = f"{url}/version"
async with session.get(url_version, headers=headers) as response:
async with session.get(url_version) as response:
if HTTPStatus(response.status) == HTTPStatus.OK:
text = await response.text()
logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
# Get available models
url_models = f"{url}/v1/models"
async with session.get(url_models, headers=headers) as response:
async with session.get(url_models) as response:
if HTTPStatus(response.status) == HTTPStatus.OK:
text = await response.text()
logger.info(f"{Color.BLUE}Models:{Color.RESET}")
@@ -1358,22 +1323,6 @@ async def main() -> None:
help="Base URL for the LLM API server",
)
parser.add_argument(
"--api-key",
type=str,
default=None,
help="API key to send as an Authorization bearer token",
)
parser.add_argument(
"--header",
action="append",
type=parse_custom_header,
default=None,
metavar="KEY=VALUE",
help="Custom request header. Can be specified multiple times. "
"Accepts 'Header-Name: value' or 'Header-Name=value'.",
)
parser.add_argument(
"-p",
"--num-clients",
@@ -1488,22 +1437,6 @@ async def main() -> None:
help="Disable stream/streaming mode (set 'stream' to False in the API request)",
)
parser.add_argument(
"--send-conversation-id",
default=False,
action="store_true",
help=(
"Inject a `conversation_id` field into each Chat Completions "
"payload. This is a non-standard OpenAI extension consumed by "
"vLLM's disaggregated multi-turn proxy "
"(examples/disaggregated/disaggregated_serving/"
"disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. "
"Leave disabled (default) when targeting strict "
"OpenAI-compatible endpoints; enable when benchmarking the "
"disaggregated proxy."
),
)
parser.add_argument(
"-e",
"--excel-output",
@@ -1592,8 +1525,7 @@ async def main() -> None:
args.model, trust_remote_code=args.trust_remote_code
)
headers = build_request_headers(args.api_key, args.header)
await get_server_info(args.url, headers=headers)
await get_server_info(args.url)
# Load the input file (either conversations of configuration file)
logger.info(f"Reading input file: {args.input_file}")
+15 -4
View File
@@ -1,5 +1,5 @@
#!/bin/bash
# Build vLLM Rust artifacts and install them into the vllm package.
# Build the vllm-rs Rust frontend binary and install it into the vllm package.
# Usage: ./build_rust.sh [--debug]
#
# By default builds in release mode. Pass --debug for faster compile times
@@ -8,6 +8,8 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
RUST_DIR="$REPO_ROOT/rust"
TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -25,9 +27,18 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
fi
if [[ "${1:-}" == "--debug" ]]; then
PROFILE_ARG="--debug"
PROFILE_ARGS=()
PROFILE_DIR="debug"
else
PROFILE_ARG="--release"
PROFILE_ARGS=(--release)
PROFILE_DIR="release"
fi
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
--manifest-path "$RUST_DIR/Cargo.toml" \
--bin vllm-rs \
--features native-tls-vendored
mkdir -p "$(dirname "$TARGET_PATH")"
cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
echo "Installed vllm-rs to $TARGET_PATH"
+3 -12
View File
@@ -166,10 +166,6 @@ elseif (S390_FOUND)
"-mtune=native")
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
message(STATUS "RISC-V detected")
if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
message(FATAL_ERROR
"VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
endif()
# VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
# by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
if(NOT DEFINED VLLM_RVV_VLEN)
@@ -193,7 +189,8 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
"RISC-V RVV is available but VLEN could not be auto-detected. "
"Please specify VLEN explicitly:\n"
" -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
" -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)")
" -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n"
" -DVLLM_RVV_VLEN=0 (force scalar, no RVV)")
endif()
endif()
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -222,7 +219,7 @@ endif()
# Build oneDNN for GEMM kernels
if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND)
if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
# Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
# TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -438,12 +435,6 @@ if(USE_ONEDNN)
${VLLM_EXT_SRC})
endif()
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
set(VLLM_EXT_SRC
"csrc/cpu/sgl-kernels/gemm_int4.cpp"
${VLLM_EXT_SRC})
endif()
if (ENABLE_X86_ISA)
set(VLLM_EXT_SRC_SGL
"csrc/cpu/sgl-kernels/conv.cpp"
-48
View File
@@ -1,48 +0,0 @@
include(FetchContent)
# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory
# instead of downloading. This is useful for local MSA development.
if(DEFINED ENV{FMHA_SM100_SRC_DIR})
set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR})
endif()
if(FMHA_SM100_SRC_DIR)
FetchContent_Declare(
fmha_sm100
SOURCE_DIR ${FMHA_SM100_SRC_DIR}
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
)
else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
)
endif()
FetchContent_GetProperties(fmha_sm100)
if(NOT fmha_sm100_POPULATED)
FetchContent_Populate(fmha_sm100)
endif()
message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
add_custom_target(fmha_sm100)
set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
install(FILES
"${FMHA_SM100_PY_ROOT}/__init__.py"
"${FMHA_SM100_PY_ROOT}/sparse.py"
DESTINATION vllm/third_party/fmha_sm100
COMPONENT fmha_sm100)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
DESTINATION vllm/third_party/fmha_sm100/cute
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE
PATTERN ".git*" EXCLUDE)
+16 -52
View File
@@ -32,35 +32,22 @@ endif()
message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_ARCHS "10.0f;12.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}")
endif()
# QUTLASS uses TARGET_CUDA_ARCH as a single preprocessor selector for all its
# sources. Do not compile a mixed SM100/SM120 arch list with one selector; prefer
# SM100 when both families are requested because that is the primary deployed
# target for this extension today.
if(QUTLASS_SM100_ARCHS)
set(QUTLASS_ARCHS "${QUTLASS_SM100_ARCHS}")
set(QUTLASS_TARGET_CC 100)
if(QUTLASS_SM120_ARCHS)
message(WARNING
"[QUTLASS] Both SM100 and SM120 archs were requested; selecting SM100 "
"because TARGET_CUDA_ARCH is a single compile-time selector.")
endif()
elseif(QUTLASS_SM120_ARCHS)
set(QUTLASS_ARCHS "${QUTLASS_SM120_ARCHS}")
set(QUTLASS_TARGET_CC 120)
else()
set(QUTLASS_ARCHS)
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;12.1a;10.0a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)")
set(QUTLASS_TARGET_CC 100)
elseif(QUTLASS_ARCHS MATCHES "12\\.[01][af]?")
set(QUTLASS_TARGET_CC 120)
else()
message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.")
endif()
set(QUTLASS_SOURCES
csrc/qutlass_registration.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -79,19 +66,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
else()
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
endif()
endif()
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
else()
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -103,23 +79,12 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${QUTLASS_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
USE_SABI 3
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
target_sources(_C PRIVATE ${QUTLASS_SOURCES})
target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
target_compile_definitions(_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -134,5 +99,4 @@ else()
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
"CUDA_ARCHS='${CUDA_ARCHS}'.")
endif()
add_custom_target(_qutlass_C)
endif()
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee
GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+2 -2
View File
@@ -487,9 +487,9 @@ endfunction()
function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.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}")
cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
endif()
set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
endfunction()
+25 -25
View File
@@ -11,25 +11,13 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
return cpu_attention::Fp8KVCacheDataType::kAuto;
}
bool cpu_attn_has_isa(const std::string& isa) {
if (isa == "rvv") {
#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
return true;
#else
return false;
#endif
}
return false;
}
torch::Tensor get_scheduler_metadata(
const int64_t num_req, const int64_t num_heads_q,
const int64_t num_heads_kv, const int64_t head_dim,
const torch::Tensor& seq_lens, at::ScalarType dtype,
const torch::Tensor& query_start_loc, const bool causal,
const torch::Tensor& query_start_loc, const bool casual,
const int64_t window_size, const std::string& isa_hint,
const bool enable_kv_split,
const std::optional<torch::Tensor>& dynamic_causal) {
const bool enable_kv_split) {
cpu_attention::ISA isa;
if (isa_hint == "amx") {
isa = cpu_attention::ISA::AMX;
@@ -56,13 +44,24 @@ torch::Tensor get_scheduler_metadata(
input.head_dim = head_dim;
input.query_start_loc = query_start_loc.data_ptr<int32_t>();
input.seq_lens = seq_lens.data_ptr<int32_t>();
input.sliding_window_size = window_size;
input.causal = causal;
if (window_size != -1) {
input.left_sliding_window_size = window_size - 1;
if (casual) {
input.right_sliding_window_size = 0;
} else {
input.right_sliding_window_size = window_size - 1;
}
} else {
input.left_sliding_window_size = -1;
if (casual) {
input.right_sliding_window_size = 0;
} else {
input.right_sliding_window_size = -1;
}
}
input.casual = casual;
input.isa = isa;
input.enable_kv_split = enable_kv_split;
input.dynamic_causal =
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
@@ -176,11 +175,10 @@ void cpu_attention_with_kv_cache(
const torch::Tensor& seq_lens, // [num_tokens]
const double scale, const bool causal,
const std::optional<torch::Tensor>& alibi_slopes, // [num_heads]
const int64_t sliding_window,
const int64_t sliding_window_left, const int64_t sliding_window_right,
const torch::Tensor& block_table, // [num_tokens, max_block_num]
const double softcap, const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux, // [num_heads]
const std::optional<torch::Tensor>& dynamic_causal, // [num_reqs]
const std::optional<torch::Tensor>& s_aux, // [num_heads]
const double k_scale = 1.0, const double v_scale = 1.0,
const std::string& kv_cache_dtype = "auto") {
TORCH_CHECK_EQ(query.dim(), 3);
@@ -222,11 +220,13 @@ void cpu_attention_with_kv_cache(
input.alibi_slopes =
alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
input.dynamic_causal =
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
input.scale = scale;
input.causal = causal;
input.sliding_window_size = sliding_window;
input.sliding_window_left = sliding_window_left;
input.sliding_window_right = sliding_window_right;
if (input.causal) {
input.sliding_window_right = 0;
}
input.softcap = static_cast<float>(softcap);
if (is_fp8) {
+34 -70
View File
@@ -388,13 +388,13 @@ class AttentionScheduler {
int32_t head_dim;
int32_t* query_start_loc;
int32_t* seq_lens;
int32_t sliding_window_size;
bool causal;
int32_t left_sliding_window_size;
int32_t right_sliding_window_size;
bool casual;
cpu_attention::ISA isa;
int32_t max_num_q_per_iter; // max Q head num can be hold in registers
int32_t kv_block_alignment; // context length alignment requirement
bool enable_kv_split;
bool* dynamic_causal;
};
static constexpr int32_t MaxQTileIterNum = 128;
@@ -403,8 +403,7 @@ class AttentionScheduler {
: available_cache_size_(cpu_utils::get_available_l2_size()) {}
torch::Tensor schedule(const ScheduleInput& input) const {
const bool causal = input.causal;
const bool is_dynamic_causal = input.dynamic_causal != nullptr;
const bool casual = input.casual;
const int32_t thread_num = omp_get_max_threads();
const int64_t cache_size = cpu_utils::get_available_l2_size();
const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
@@ -435,7 +434,8 @@ class AttentionScheduler {
const int32_t default_tile_token_num = default_tile_size / q_head_per_kv;
const int32_t split_kv_q_token_num_threshold =
input.enable_kv_split ? 1 : 0;
const int32_t sliding_window_size = input.sliding_window_size;
const int32_t left_sliding_window_size = input.left_sliding_window_size;
const int32_t right_sliding_window_size = input.right_sliding_window_size;
TORCH_CHECK_LE(split_kv_q_token_num_threshold * q_head_per_kv, 16);
// get total kv len
@@ -444,9 +444,7 @@ class AttentionScheduler {
const int32_t seq_len = input.seq_lens[req_id];
const int32_t q_token_num =
input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
const bool req_causal =
is_dynamic_causal ? input.dynamic_causal[req_id] : causal;
const int32_t q_start_pos = seq_len - q_token_num;
const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
const int32_t kv_start_pos = 0;
const int32_t kv_end_pos = seq_len;
@@ -458,7 +456,7 @@ class AttentionScheduler {
const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
sliding_window_size, req_causal);
left_sliding_window_size, right_sliding_window_size);
const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
kv_len_alignment);
@@ -486,9 +484,7 @@ class AttentionScheduler {
const int32_t seq_len = input.seq_lens[req_id];
const int32_t q_token_num =
input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
const bool req_causal =
is_dynamic_causal ? input.dynamic_causal[req_id] : causal;
const int32_t q_start_pos = seq_len - q_token_num;
const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
const int32_t kv_start_pos = 0;
const int32_t kv_end_pos = seq_len;
int32_t local_split_id = 0;
@@ -502,7 +498,7 @@ class AttentionScheduler {
const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
sliding_window_size, req_causal);
left_sliding_window_size, right_sliding_window_size);
const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
kv_len_alignment);
@@ -712,41 +708,15 @@ class AttentionScheduler {
return metadata_tensor;
}
FORCE_INLINE static std::pair<int32_t, int32_t> calcu_sliding_window_size(
int32_t window_size, bool causal) {
int32_t left_sliding_window_size, right_sliding_window_size;
if (window_size != -1) {
left_sliding_window_size = window_size - 1;
if (causal) {
right_sliding_window_size = 0;
} else {
right_sliding_window_size = window_size - 1;
}
} else {
left_sliding_window_size = -1;
if (causal) {
right_sliding_window_size = 0;
} else {
right_sliding_window_size = -1;
}
}
return {left_sliding_window_size, right_sliding_window_size};
}
FORCE_INLINE static std::pair<int32_t, int32_t> calcu_kv_tile_pos(
int32_t kv_left_pos, int32_t kv_right_pos, int32_t q_left_pos,
int32_t q_right_pos, int32_t window_size, bool causal) {
auto [left_sliding_window_size, right_sliding_window_size] =
calcu_sliding_window_size(window_size, causal);
if (left_sliding_window_size != -1) {
kv_left_pos =
std::max(kv_left_pos, q_left_pos - left_sliding_window_size);
int32_t q_right_pos, int32_t sliding_window_left,
int32_t sliding_window_right) {
if (sliding_window_left != -1) {
kv_left_pos = std::max(kv_left_pos, q_left_pos - sliding_window_left);
}
if (right_sliding_window_size != -1) {
kv_right_pos =
std::min(kv_right_pos, q_right_pos + right_sliding_window_size);
if (sliding_window_right != -1) {
kv_right_pos = std::min(kv_right_pos, q_right_pos + sliding_window_right);
}
return {kv_left_pos, kv_right_pos};
}
@@ -835,10 +805,10 @@ struct AttentionInput {
int32_t* block_table;
float* alibi_slopes;
c10::BFloat16* s_aux;
bool* dynamic_causal;
float scale;
bool causal;
int32_t sliding_window_size;
int32_t sliding_window_left;
int32_t sliding_window_right;
float softcap;
// FP8 KV cache scales (used by FP8 attention implementations)
float k_scale_fp8 = 1.0f;
@@ -852,8 +822,8 @@ struct AttentionInput {
logits_buffer_t *__restrict__ logits_buffer, \
float *__restrict__ partial_q_buffer, float *__restrict__ max_buffer, \
float *__restrict__ sum_buffer, int32_t *__restrict__ block_table, \
const int32_t kv_end_pos, const int32_t kv_tile_start_pos, \
const int32_t kv_tile_end_pos, const int32_t kv_tile_token_num, \
const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos, \
const int32_t kv_tile_token_num, \
const int64_t kv_cache_num_blocks_stride, const int32_t q_head_num, \
const int32_t q_token_num, const int32_t q_tile_start_pos, \
const int32_t q_heads_per_kv, const int32_t block_size, \
@@ -864,7 +834,7 @@ struct AttentionInput {
#define CPU_ATTENTION_PARAMS \
q_heads_buffer, k_head_cache_ptr, v_head_cache_ptr, logits_buffer, \
partial_q_buffer, max_buffer, sum_buffer, block_table, kv_end_pos, \
partial_q_buffer, max_buffer, sum_buffer, block_table, \
kv_tile_start_pos, kv_tile_end_pos, kv_tile_token_num, \
kv_cache_num_blocks_stride, q_head_num, q_token_num, q_tile_start_pos, \
q_heads_per_kv, block_size, left_window_size, right_window_size, scale, \
@@ -947,7 +917,6 @@ class AttentionMainLoop {
// - max_buffer: [MaxQHeadNumPerIteration, 1], store max logits
// - sum_buffer: [MaxQHeadNumPerIteration, 1], store sum of exp
// - block_table
// - kv_end_pos: un-aligned end position of KV cache
// - kv_tile_start_pos: start position of KV cache, aligned to
// BlockSizeAlignment
// - kv_tile_end_pos: end position of KV cache, aligned to
@@ -1074,7 +1043,7 @@ class AttentionMainLoop {
}
apply_mask(logits_buffer, kv_tile_token_num, q_tile_start_pos,
kv_end_pos, kv_tile_start_pos, kv_tile_end_pos, q_token_num,
kv_tile_start_pos, kv_tile_end_pos, q_token_num,
q_heads_per_kv, left_window_size, right_window_size);
// if (debug_info){
@@ -1157,7 +1126,7 @@ class AttentionMainLoop {
void apply_mask(logits_buffer_t* __restrict__ logits_buffer,
const int64_t logits_buffer_stride,
const int32_t q_tile_start_pos, const int32_t kv_end_pos,
const int32_t q_tile_start_pos,
const int32_t kv_tile_start_pos,
const int32_t kv_tile_end_pos, const int32_t q_token_num,
const int32_t q_heads_per_kv,
@@ -1185,7 +1154,7 @@ class AttentionMainLoop {
std::max(kv_tile_start_pos,
curr_token_pos + sliding_window_right + 1));
}
return std::min(pos, kv_end_pos);
return pos;
}();
int32_t left_invalid_token_num = left_kv_pos - kv_tile_start_pos;
@@ -1472,16 +1441,15 @@ class AttentionMainLoop {
const int64_t q_head_num_stride = input->query_num_heads_stride;
const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
const int64_t kv_cache_block_num_stride = input->cache_num_blocks_stride;
const int32_t sliding_window_size = input->sliding_window_size;
const int32_t sliding_window_left = input->sliding_window_left;
const int32_t sliding_window_right = input->sliding_window_right;
const int32_t block_size = input->block_size;
const float scale = input->scale;
const float softcap_scale = input->softcap;
const float* alibi_slopes = input->alibi_slopes;
const c10::BFloat16* s_aux = input->s_aux;
const bool* dynamic_causal = input->dynamic_causal;
const bool is_dynamic_causal = dynamic_causal != nullptr;
const bool causal = input->causal;
const bool casual = input->causal;
int32_t* const block_table = input->block_table;
const int64_t block_table_stride = input->blt_num_tokens_stride;
@@ -1564,11 +1532,6 @@ class AttentionMainLoop {
&curr_workitem_groups[workitem_group_idx];
const int32_t current_group_idx = current_workitem_group->req_id;
const int32_t current_group_causal =
is_dynamic_causal ? dynamic_causal[current_group_idx] : causal;
auto [sliding_window_left, sliding_window_right] =
AttentionScheduler::calcu_sliding_window_size(
sliding_window_size, current_group_causal);
const int32_t kv_start_pos =
current_workitem_group->kv_split_pos_start;
const int32_t kv_end_pos = current_workitem_group->kv_split_pos_end;
@@ -1596,7 +1559,8 @@ class AttentionMainLoop {
const int32_t q_end = input->query_start_loc[current_group_idx + 1];
const int32_t q_start = input->query_start_loc[current_group_idx];
const int32_t seq_len = input->seq_lens[current_group_idx];
const int32_t q_start_pos = seq_len - (q_end - q_start);
const int32_t q_start_pos =
(casual ? seq_len - (q_end - q_start) : 0);
const int32_t block_num = (seq_len + block_size - 1) / block_size;
// Only apply sink for the first KV split
bool use_sink = (s_aux != nullptr &&
@@ -1646,8 +1610,8 @@ class AttentionMainLoop {
const auto [kv_tile_start_pos, kv_tile_end_pos] =
AttentionScheduler::calcu_kv_tile_pos(
kv_start_pos, kv_end_pos, q_tile_start_pos,
q_tile_end_pos, sliding_window_size,
current_group_causal);
q_tile_end_pos, sliding_window_left,
sliding_window_right);
const auto [rounded_kv_tile_start_pos, rounded_kv_tile_end_pos] =
AttentionScheduler::align_kv_tile_pos(
kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
@@ -1760,8 +1724,8 @@ class AttentionMainLoop {
actual_kv_tile_pos_right] =
AttentionScheduler::calcu_kv_tile_pos(
kv_tile_pos_left, kv_tile_pos_right, q_tile_pos_left,
q_tile_pos_right, sliding_window_size,
current_group_causal);
q_tile_pos_right, sliding_window_left,
sliding_window_right);
const int32_t q_iter_idx =
q_head_tile_token_offset / curr_max_q_token_num_per_iter;
@@ -1825,7 +1789,7 @@ class AttentionMainLoop {
attn_impl.template execute_attention<Attention>(
curr_q_heads_buffer, curr_k_cache, curr_v_cache,
logits_buffer, curr_partial_q_buffer, curr_max_buffer,
curr_sum_buffer, curr_block_table, kv_end_pos,
curr_sum_buffer, curr_block_table,
aligned_actual_kv_tile_pos_left,
aligned_actual_kv_tile_pos_right, actual_kv_token_num,
kv_cache_block_num_stride, q_tile_head_num,
+7 -5
View File
@@ -1,5 +1,3 @@
#include <sleef.h>
#include "cpu/cpu_types.hpp"
#include "cpu/utils.hpp"
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
@@ -165,6 +163,7 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
vec_op::FP32Vec16 w2_vec(0.5);
vec_op::FP32Vec16 w3_vec(0.044715);
alignas(64) float temp[16];
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
@@ -172,9 +171,12 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
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
vec_op::FP32Vec16 tanh_vec(Sleef_tanhf16_u10(inner_vec.reg));
inner_vec.save(temp);
for (int32_t i = 0; i < 16; ++i) {
temp[i] = std::tanh(temp[i]);
}
vec_op::FP32Vec16 tanh_vec(temp);
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);
-4
View File
@@ -57,10 +57,6 @@ typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
__attribute__((riscv_rvv_vector_bits(1024)));
// int8
typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t
__attribute__((riscv_rvv_vector_bits(128)));
// int32
typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
__attribute__((riscv_rvv_vector_bits(256)));
+39 -49
View File
@@ -9,14 +9,10 @@
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <torch/all.h>
#include "float_convert.hpp"
namespace vec_op {
// FP8 KV cache is not supported on RISC-V. These tag types and the
@@ -249,7 +245,8 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[8];
for (int i = 0; i < 8; ++i) {
tmp[i] = bf16_to_float(u16[i]);
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
}
@@ -259,7 +256,9 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 8; ++i) {
u16[i] = float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
@@ -267,7 +266,9 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
u16[i] = float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -276,8 +277,10 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 8; ++i) {
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
@@ -289,7 +292,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[16];
for (int i = 0; i < 16; ++i) {
tmp[i] = bf16_to_float(u16[i]);
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
}
@@ -302,7 +306,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 16; ++i) {
u16[i] = float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
@@ -310,7 +316,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
u16[i] = float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -319,8 +327,10 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 16; ++i) {
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
@@ -333,7 +343,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[32];
for (int i = 0; i < 32; ++i) {
tmp[i] = bf16_to_float(u16[i]);
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
}
@@ -360,7 +371,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 32; ++i) {
u16[i] = float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
@@ -369,7 +382,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
u16[i] = float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
@@ -379,8 +394,10 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 32; ++i) {
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
float_to_bf16(tmp[i]);
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
@@ -717,18 +734,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
return FP32Vec16(
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
return FP32Vec16(
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, elem_num));
}
FP32Vec16 min(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
return FP32Vec16(
RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, elem_num));
}
FP32Vec16 abs() const {
return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
}
@@ -858,27 +867,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
}
};
struct INT8Vec16 : public Vec<INT8Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_i8x16_t reg;
explicit INT8Vec16(const FP32Vec16& vec) {
auto i32_vec =
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(vec.reg, VEC_ELEM_NUM);
auto i16_vec = RVVI(__riscv_vnclip_wx_i16, LMUL_256)(
i32_vec, 0, __RISCV_VXRM_RNU, VEC_ELEM_NUM);
reg = RVVI(__riscv_vnclip_wx_i8, LMUL_128)(i16_vec, 0, __RISCV_VXRM_RNU,
VEC_ELEM_NUM);
}
void save(int8_t* ptr) const {
RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
}
void save(int8_t* ptr, int elem_num) const {
RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, elem_num);
}
};
// ============================================================================
// Type Traits & Global Helpers
// ============================================================================
@@ -968,7 +956,9 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
#else
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
*reinterpret_cast<uint16_t*>(ptr) = float_to_bf16(v);
uint32_t val;
std::memcpy(&val, &v, 4);
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
}
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
+2 -3
View File
@@ -3,9 +3,7 @@
#define CPU_TYPES_VXE_HPP
#include <vecintrin.h>
#include <bit>
#include <cmath>
#include <cstdint>
#include <limits>
#include <torch/all.h>
namespace vec_op {
@@ -819,7 +817,8 @@ inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
// intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
// produce incorrect results for some inputs. Process each of the 4 vectors
// separately.
uint32_t in = std::bit_cast<uint32_t>(v);
uint32_t in;
std::memcpy(&in, &v, sizeof(in));
uint32_t s = (in & 0x80000000) >> 16; // Sign
uint32_t e = (in & 0x7F800000) >> 23; // Exponent
-39
View File
@@ -4,9 +4,6 @@
#ifdef CPU_CAPABILITY_AMXBF16
#include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
#endif
#if defined(__riscv_v)
#include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
#endif
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
@@ -322,8 +319,6 @@ void cpu_gemm_wna16(
return ISA::AMX;
} else if (isa_hint == "vec") {
return ISA::VEC;
} else if (isa_hint == "rvv") {
return ISA::RVV;
} else {
TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
}
@@ -402,40 +397,6 @@ void cpu_gemm_wna16(
pack_factor);
return;
}
} else if (isa == ISA::RVV) {
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
if (has_zp) {
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
scales_group_stride, zeros_group_stride, group_num, group_size,
pack_factor);
return;
}
if (use_desc_act) {
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
scales_group_stride, zeros_group_stride, group_num, group_size,
pack_factor);
return;
} else {
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
scales_group_stride, zeros_group_stride, group_num, group_size,
pack_factor);
return;
}
}
});
}
+15 -13
View File
@@ -1,15 +1,14 @@
#pragma once
#include <bit>
#include <cstdint>
inline float bf16_to_float(uint16_t bf16) {
static float bf16_to_float(uint16_t bf16) {
uint32_t bits = static_cast<uint32_t>(bf16) << 16;
return std::bit_cast<float>(bits);
float fp32;
std::memcpy(&fp32, &bits, sizeof(fp32));
return fp32;
}
inline uint16_t float_to_bf16(float fp32) {
uint32_t bits = std::bit_cast<uint32_t>(fp32);
static uint16_t float_to_bf16(float fp32) {
uint32_t bits;
std::memcpy(&bits, &fp32, sizeof(fp32));
return static_cast<uint16_t>(bits >> 16);
}
@@ -19,13 +18,14 @@ inline uint16_t float_to_bf16(float fp32) {
* Codes below copied from
* https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
*************************************************/
inline uint16_t float_to_fp16(float fp32) {
static uint16_t float_to_fp16(float fp32) {
uint16_t fp16;
unsigned x;
unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
unsigned sign, exponent, mantissa;
uint32_t x = std::bit_cast<uint32_t>(fp32);
std::memcpy(&x, &fp32, sizeof(fp32));
u = (x & 0x7fffffff);
// Get rid of +NaN/-NaN case first.
@@ -77,11 +77,12 @@ inline uint16_t float_to_fp16(float fp32) {
return fp16;
}
inline float fp16_to_float(uint16_t fp16) {
static float fp16_to_float(uint16_t fp16) {
unsigned sign = ((fp16 >> 15) & 1);
unsigned exponent = ((fp16 >> 10) & 0x1f);
unsigned mantissa = ((fp16 & 0x3ff) << 13);
uint32_t temp;
int temp;
float fp32;
if (exponent == 0x1f) { /* NaN or Inf */
mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
exponent = 0xff;
@@ -100,5 +101,6 @@ inline float fp16_to_float(uint16_t fp16) {
exponent += 0x70;
}
temp = ((sign << 31) | (exponent << 23) | mantissa);
return std::bit_cast<float>(temp);
std::memcpy(&fp32, &temp, sizeof(temp));
return fp32;
}
+1 -1
View File
@@ -11,7 +11,7 @@ import os
HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
# Head dimensions divisible by 16 but not 32 (VEC16 only)
HEAD_DIMS_16 = [48, 80, 112]
HEAD_DIMS_16 = [80, 112]
# ISA types
ISA_TYPES = {
+22 -41
View File
@@ -4,9 +4,8 @@ namespace {
template <typename scalar_t>
void rms_norm_impl(scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ weight, const bool has_weight,
const float epsilon, const int num_tokens,
const int hidden_size) {
const scalar_t* __restrict__ weight, const float epsilon,
const int num_tokens, const int hidden_size) {
using scalar_vec_t = vec_op::vec_t<scalar_t>;
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
@@ -28,15 +27,12 @@ void rms_norm_impl(scalar_t* __restrict__ out,
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
scalar_vec_t x(input_p + j);
scalar_vec_t w(weight + j);
vec_op::FP32Vec8 fp32_x(x);
vec_op::FP32Vec8 fp32_out;
if (has_weight) {
scalar_vec_t w(weight + j);
vec_op::FP32Vec8 fp32_w(w);
fp32_out = fp32_x * fp32_s_variance * fp32_w;
} else {
fp32_out = fp32_x * fp32_s_variance;
}
vec_op::FP32Vec8 fp32_w(w);
vec_op::FP32Vec8 fp32_out = fp32_x * fp32_s_variance * fp32_w;
scalar_vec_t out(fp32_out);
out.save(output_p + j);
@@ -48,8 +44,8 @@ template <typename scalar_t>
void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
scalar_t* __restrict__ residual,
const scalar_t* __restrict__ weight,
const bool has_weight, const float epsilon,
const int num_tokens, const int hidden_size) {
const float epsilon, const int num_tokens,
const int hidden_size) {
using scalar_vec_t = vec_op::vec_t<scalar_t>;
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
@@ -76,18 +72,13 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
vec_op::FP32Vec8 fp32_s_variance(s_variance);
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
vec_op::FP32Vec8 fp32_out;
if (has_weight) {
scalar_vec_t w(weight + j);
scalar_vec_t res(residual_p + j);
vec_op::FP32Vec8 fp32_w(w);
vec_op::FP32Vec8 fp32_res(res);
fp32_out = fp32_res * fp32_s_variance * fp32_w;
} else {
scalar_vec_t res(residual_p + j);
vec_op::FP32Vec8 fp32_res(res);
fp32_out = fp32_res * fp32_s_variance;
}
scalar_vec_t w(weight + j);
scalar_vec_t res(residual_p + j);
vec_op::FP32Vec8 fp32_w(w);
vec_op::FP32Vec8 fp32_res(res);
vec_op::FP32Vec8 fp32_out = fp32_res * fp32_s_variance * fp32_w;
scalar_vec_t out(fp32_out);
out.save(input_p + j);
@@ -96,41 +87,31 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
}
} // namespace
void rms_norm(torch::Tensor& out, torch::Tensor& input,
std::optional<torch::Tensor> weight, double epsilon) {
void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
double epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
const bool has_weight = weight.has_value();
if (has_weight) {
TORCH_CHECK(weight->is_contiguous());
}
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] {
CPU_KERNEL_GUARD_IN(rms_norm_impl)
rms_norm_impl(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
has_weight ? weight->data_ptr<scalar_t>() : nullptr,
has_weight, epsilon, num_tokens, hidden_size);
weight.data_ptr<scalar_t>(), epsilon, num_tokens,
hidden_size);
CPU_KERNEL_GUARD_OUT(rms_norm_impl)
});
}
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
std::optional<torch::Tensor> weight, double epsilon) {
torch::Tensor& weight, double epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
const bool has_weight = weight.has_value();
if (has_weight) {
TORCH_CHECK(weight->scalar_type() == input.scalar_type());
TORCH_CHECK(weight->is_contiguous());
}
VLLM_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "fused_add_rms_norm_impl", [&] {
CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl)
fused_add_rms_norm_impl(
input.data_ptr<scalar_t>(), residual.data_ptr<scalar_t>(),
has_weight ? weight->data_ptr<scalar_t>() : nullptr, has_weight,
epsilon, num_tokens, hidden_size);
weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl)
});
}
-228
View File
@@ -1,228 +0,0 @@
#ifndef CPU_MICRO_GEMM_RVV_HPP
#define CPU_MICRO_GEMM_RVV_HPP
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
#if defined(__riscv_v)
namespace cpu_micro_gemm {
namespace {
constexpr int32_t RVV_MGEMM_N8 = 8;
constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
template <typename scalar_t>
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
template <>
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
}
template <>
FORCE_INLINE fixed_fp32x8_t
load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
#if defined(__riscv_zvfh)
fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
#else
alignas(32) float values[RVV_MGEMM_N8];
for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
values[i] = static_cast<float>(ptr[i]);
}
return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
#endif
}
template <>
FORCE_INLINE fixed_fp32x8_t
load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
#if defined(__riscv_zvfbfmin)
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
fixed_bf16x8_t vec =
RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
#else
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
#endif
}
// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
template <int32_t M, typename scalar_t>
FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
const scalar_t* __restrict__ b_ptr,
float* __restrict__ c_ptr,
const int64_t lda,
const int64_t ldc, const int32_t k,
const bool accum_c) {
static_assert(0 < M && M <= 8);
#define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
#define RVV_IF_M(i) if constexpr (M > (i))
#define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
RVV_ROWS_APPLY(RVV_DECL_A)
#undef RVV_DECL_A
#define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
RVV_ROWS_APPLY(RVV_DECL_ACC)
#undef RVV_DECL_ACC
#define RVV_INIT_ACC(i) \
RVV_IF_M(i) { \
if (accum_c) { \
acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \
RVV_MGEMM_N8); \
} else { \
acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
} \
}
RVV_ROWS_APPLY(RVV_INIT_ACC)
#undef RVV_INIT_ACC
int32_t k_idx = 0;
for (; k_idx + 3 < k; k_idx += 4) {
#define RVV_FMA_ROW(i, K_OFFSET) \
RVV_IF_M(i) { \
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
RVV_MGEMM_N8); \
}
#define RVV_STEP_K(K_OFFSET) \
{ \
fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>( \
b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
RVV_FMA_ROW(0, K_OFFSET) \
RVV_FMA_ROW(1, K_OFFSET) \
RVV_FMA_ROW(2, K_OFFSET) \
RVV_FMA_ROW(3, K_OFFSET) \
RVV_FMA_ROW(4, K_OFFSET) \
RVV_FMA_ROW(5, K_OFFSET) \
RVV_FMA_ROW(6, K_OFFSET) \
RVV_FMA_ROW(7, K_OFFSET) \
}
RVV_STEP_K(0)
RVV_STEP_K(1)
RVV_STEP_K(2)
RVV_STEP_K(3)
#undef RVV_STEP_K
#undef RVV_FMA_ROW
}
for (; k_idx < k; ++k_idx) {
fixed_fp32x8_t b =
load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
#define RVV_TAIL_ROW(i) \
RVV_IF_M(i) { \
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
}
RVV_ROWS_APPLY(RVV_TAIL_ROW)
#undef RVV_TAIL_ROW
}
#define RVV_STORE_ROW(i) \
RVV_IF_M(i) { \
RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
RVV_MGEMM_N8); \
}
RVV_ROWS_APPLY(RVV_STORE_ROW)
#undef RVV_STORE_ROW
#undef RVV_ROWS_APPLY
#undef RVV_IF_M
}
template <int32_t M, typename scalar_t>
FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
static_assert(0 < M && M <= 8);
scalar_t* __restrict__ curr_b_0 = b_ptr;
scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
accum_c);
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
ldc, k, accum_c);
}
class TileGemmRVV {
public:
template <typename scalar_t>
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
switch (m) {
case 1:
gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
break;
case 2:
gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
break;
case 3:
gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
break;
case 4:
gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
break;
case 5:
gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
break;
case 6:
gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
break;
case 7:
gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
break;
case 8:
gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
break;
}
}
};
} // namespace
template <typename scalar_t>
class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
public:
static constexpr int32_t MaxMSize = 8;
static constexpr int32_t NSize = 32;
public:
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
}
static void pack_weight(const scalar_t* __restrict__ weight,
scalar_t* __restrict__ packed_weight,
const int32_t output_size, const int32_t input_size) {
TORCH_CHECK_EQ(output_size % 16, 0);
for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
scalar_t* __restrict__ curr_packed_weight =
packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
*curr_packed_weight = *curr_weight;
curr_packed_weight += 16;
++curr_weight;
}
}
}
};
} // namespace cpu_micro_gemm
#endif // defined(__riscv_v)
#endif // CPU_MICRO_GEMM_RVV_HPP
+1 -37
View File
@@ -268,23 +268,6 @@ void _dequant_gemm_accum_small_M(
_dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
#endif
template <int64_t N, int64_t ldb>
inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t n) {
// B is packed as [_block_k / 4, N / 2, 4] for VNNI4. Each byte stores two
// columns from adjacent 8-column groups for one K lane.
constexpr int64_t n_group_size = 8;
constexpr int64_t vnni_size = 4;
static_assert(N % (2 * n_group_size) == 0);
int64_t n_group = n / n_group_size;
int64_t ni = n % n_group_size;
int64_t ki = k % vnni_size;
int64_t k_base = k - ki;
int64_t packed_n = (n_group / 2) * n_group_size + ni;
uint8_t packed = B[k_base * ldb + packed_n * vnni_size + ki];
return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
}
template <int64_t N, int64_t ldb, bool sym_quant_act>
void _dequant_gemm_accum(
float* C,
@@ -338,24 +321,7 @@ void _dequant_gemm_accum(
} else
#endif
{
for (int64_t m = 0; m < M; ++m) {
for (int64_t n = 0; n < N; ++n) {
int32_t acc = 0;
for (int64_t k = 0; k < K; ++k) {
int32_t b = load_uint4_vnni<N, ldb>(B, k, n) - qzeros_b[n];
if constexpr (sym_quant_act) {
const int8_t* A_s8 = reinterpret_cast<const int8_t*>(A);
acc += static_cast<int32_t>(A_s8[m * lda + k]) * b;
} else {
acc += static_cast<int32_t>(A[m * lda + k]) * b;
}
}
if constexpr (!sym_quant_act) {
acc -= qzeros_a[m] * compensation[n];
}
C[m * ldc + n] += static_cast<float>(acc) * scales_a[m] * scales_b[n];
}
}
TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!");
}
}
@@ -530,11 +496,9 @@ void _da8w4_linear_impl(
store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
}
}
#if defined(CPU_CAPABILITY_AVX512)
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
#endif
});
}
+1 -1
View File
@@ -245,7 +245,7 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric
for (int64_t k = 0; k < K; ++k) {
const float val = static_cast<float>(A[k]) * inv_scale;
Aq[k] = static_cast<uint8_t>(static_cast<int32_t>(std::round(val)) + 128);
Aq[k] = (uint8_t)(std::round(val)) + 128;
}
As = scale;
}
+25 -35
View File
@@ -146,16 +146,13 @@ at::Tensor causal_conv1d_update_cpu(
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const std::string& activation);
bool cpu_attn_has_isa(const std::string& isa);
torch::Tensor get_scheduler_metadata(
const int64_t num_req, const int64_t num_heads_q,
const int64_t num_heads_kv, const int64_t head_dim,
const torch::Tensor& seq_lens, at::ScalarType dtype,
const torch::Tensor& query_start_loc, const bool casual,
const int64_t window_size, const std::string& isa_hint,
const bool enable_kv_split,
const std::optional<torch::Tensor>& dynamic_causal);
const bool enable_kv_split);
void cpu_attn_reshape_and_cache(const torch::Tensor& key,
const torch::Tensor& value,
@@ -172,10 +169,10 @@ void cpu_attention_with_kv_cache(
const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
const double scale, const bool causal,
const std::optional<torch::Tensor>& alibi_slopes,
const int64_t sliding_window_left, const torch::Tensor& block_table,
const double softcap, const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux,
const std::optional<torch::Tensor>& dynamic_causal, const double k_scale,
const int64_t sliding_window_left, const int64_t sliding_window_right,
const torch::Tensor& block_table, const double softcap,
const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux, const double k_scale,
const double v_scale, const std::string& kv_cache_dtype);
// Note: just for avoiding importing errors
@@ -312,13 +309,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Layernorm
// Apply Root Mean Square (RMS) Normalization to the input tensor.
ops.def(
"rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> "
"rms_norm(Tensor! out, Tensor input, Tensor weight, float epsilon) -> "
"()");
ops.impl("rms_norm", torch::kCPU, &rms_norm);
// In-place fused Add and RMS Normalization.
ops.def(
"fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
"fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
"float epsilon) -> ()");
ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm);
@@ -332,9 +329,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
// Quantization
#if defined(__AVX512F__) || defined(__AVX2__) || \
(defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
defined(__riscv_v)
#if defined(__AVX512F__) || defined(__AVX2__) || \
(defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
// Helper function to release oneDNN handlers
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
&release_dnnl_matmul_handler);
@@ -432,6 +428,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
&int8_scaled_mm_with_quant);
// Adapted from sglang: INT4 W4A8 kernels
ops.def(
"convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
"scales, int quant_method_4bit) -> (Tensor, "
"Tensor, Tensor)");
ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
&convert_weight_packed_scale_zp);
ops.def(
"int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
"Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
// Adapted from sglang: FP8 W8A16 kernel
ops.def(
"fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
@@ -458,23 +467,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
defined(__AVX512VNNI__)) || \
defined(__riscv)
// Adapted from sglang: INT4 W4A8 kernels
ops.def(
"convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
"scales, int quant_method_4bit) -> (Tensor, "
"Tensor, Tensor)");
ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
&convert_weight_packed_scale_zp);
ops.def(
"int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
"Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
#endif
// Adapted from sglang: GDN kernels
ops.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -499,12 +491,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
// CPU attention kernels
ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa);
ops.def(
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
"int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
"query_start_loc, bool casual, int window_size, str isa_hint, bool "
"enable_kv_split, Tensor? dynamic_causal) -> Tensor",
"enable_kv_split) -> Tensor",
&get_scheduler_metadata);
ops.def(
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
@@ -516,9 +507,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
"value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
"seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
"sliding_window_size, Tensor block_table, "
"float softcap, Tensor scheduler_metadata, Tensor? s_aux, Tensor? "
"dynamic_causal, "
"sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
"float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
"()",
&cpu_attention_with_kv_cache);
+1 -3
View File
@@ -8,15 +8,13 @@
#include "cpu/cpu_types.hpp"
namespace cpu_utils {
enum class ISA { AMX, VEC, RVV };
enum class ISA { AMX, VEC };
inline ISA get_isa(const std::string& isa) {
if (isa == "amx") {
return ISA::AMX;
} else if (isa == "vec") {
return ISA::VEC;
} else if (isa == "rvv") {
return ISA::RVV;
} else {
TORCH_CHECK(false, "Invalid isa type: " + isa);
}
+73 -12
View File
@@ -1,6 +1,7 @@
// A CUDAPluggableAllocator based on cumem* APIs.
// Important: allocation size, CUdeviceptr and CUmemGenericAllocationHandle*
// need to be unsigned long long
#include <atomic>
#include <iostream>
#include "cumem_allocator_compat.h"
@@ -116,6 +117,59 @@ void ensure_context(unsigned long long device) {
}
}
// ---------------------------------------------------------------------------
// Cached fabric handle probe (CUDA 12.4+, NVIDIA only):
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040
// Per-device cache: 0 = not probed, 1 = supported, 2 = not supported
static constexpr int MAX_DEVICES = 32;
static std::atomic<int> fabric_support[MAX_DEVICES] = {};
static bool probe_fabric_support(unsigned long long device) {
if (device >= MAX_DEVICES) return false;
int cached = fabric_support[device].load(std::memory_order_acquire);
if (cached != 0) return cached == 1;
int fab_flag = 0;
CUresult r = cuDeviceGetAttribute(
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
if (r != CUDA_SUCCESS || !fab_flag) {
fabric_support[device].store(2, std::memory_order_release);
return false;
}
// Attribute says supported — verify with a real allocation.
// cuDeviceGetAttribute can report supported even when IMEX is not
// configured, so we need a real probe.
CUmemAllocationProp probe_prop = {};
probe_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
probe_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
probe_prop.location.id = device;
probe_prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
size_t granularity;
r = cuMemGetAllocationGranularity(&granularity, &probe_prop,
CU_MEM_ALLOC_GRANULARITY_MINIMUM);
if (r != CUDA_SUCCESS) {
fabric_support[device].store(2, std::memory_order_release);
return false;
}
CUmemGenericAllocationHandle test_handle;
r = cuMemCreate(&test_handle, granularity, &probe_prop, 0);
if (r == CUDA_SUCCESS) {
cuMemRelease(test_handle);
fabric_support[device].store(1, std::memory_order_release);
return true;
}
fabric_support[device].store(2, std::memory_order_release);
return false;
}
#endif
// ---------------------------------------------------------------------------
void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
#ifndef USE_ROCM
CUmemGenericAllocationHandle* p_memHandle) {
@@ -136,32 +190,40 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
CUresult rdma_result = cuDeviceGetAttribute(
&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
device);
if (rdma_result == CUDA_SUCCESS &&
flag) { // support GPUDirect RDMA if possible
if (rdma_result == CUDA_SUCCESS && flag) {
prop.allocFlags.gpuDirectRDMACapable = 1;
}
int fab_flag = 0;
CUresult fab_result = cuDeviceGetAttribute(
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
if (fab_result == CUDA_SUCCESS &&
fab_flag) { // support fabric handle if possible
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
if (probe_fabric_support(device)) {
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
} else {
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
}
#else
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
#endif
#endif
#ifndef USE_ROCM
// Allocate memory using cuMemCreate
CUresult ret = (CUresult)cuMemCreate(p_memHandle, size, &prop, 0);
if (ret) {
if (fab_flag &&
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
// Safety net: if fabric was probed as available but this allocation
// still fails, fall back to POSIX FD and update the cache.
if (device < MAX_DEVICES &&
fabric_support[device].load(std::memory_order_acquire) == 1 &&
(ret == CUDA_ERROR_NOT_PERMITTED || ret == CUDA_ERROR_NOT_SUPPORTED)) {
// Fabric allocation may fail without multi-node nvlink,
// fallback to POSIX file descriptor
fabric_support[device].store(2, std::memory_order_release);
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0));
} else {
CUDA_CHECK(ret);
}
#else
CUDA_CHECK(ret);
#endif
}
if (error_code != 0) {
return;
@@ -326,14 +388,13 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
// first allocation, align the size, and reserve an address, and also allocate
// a CUmemGenericAllocationHandle
// Define memory allocation properties
// No handle type here; create_and_map sets fabric/POSIX as needed.
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = device;
prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE;
// Check if the allocation is supported
size_t granularity;
CUDA_CHECK(cuMemGetAllocationGranularity(&granularity, &prop,
CU_MEM_ALLOC_GRANULARITY_MINIMUM));
@@ -57,13 +57,13 @@ VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = {
}
VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = {
DataType.u8: "torch::headeronly::ScalarType::Byte",
DataType.s8: "torch::headeronly::ScalarType::Char",
DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn",
DataType.s32: "torch::headeronly::ScalarType::Int",
DataType.f16: "torch::headeronly::ScalarType::Half",
DataType.bf16: "torch::headeronly::ScalarType::BFloat16",
DataType.f32: "torch::headeronly::ScalarType::Float",
DataType.u8: "at::ScalarType::Byte",
DataType.s8: "at::ScalarType::Char",
DataType.e4m3: "at::ScalarType::Float8_e4m3fn",
DataType.s32: "at::ScalarType::Int",
DataType.f16: "at::ScalarType::Half",
DataType.bf16: "at::ScalarType::BFloat16",
DataType.f32: "at::ScalarType::Float",
}
VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = {

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