karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

Merge branch 'main' into wentao-fix-es-v2-bug

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
yewentao256
2026-06-03 18:13:22 +00:00
parent d8fe290d36 e5232679a3
commit 133ef0cbff
566 changed files with 28025 additions and 7233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/ci_config_rocm.yaml
+23
View File
@@ -0,0 +1,23 @@
name: vllm_rocm_ci
job_dirs:
- ".buildkite/hardware_tests"
run_all_patterns:
- "docker/Dockerfile.rocm"
- "docker/Dockerfile.rocm_base"
- "docker/ci-rocm.hcl"
- "docker/docker-bake-rocm.hcl"
- ".buildkite/hardware_tests/amd.yaml"
- ".buildkite/scripts/ci-bake-rocm.sh"
- ".buildkite/scripts/hardware_ci/run-amd-test.py"
- ".buildkite/scripts/hardware_ci/run-amd-test.sh"
- "CMakeLists.txt"
- "requirements/common.txt"
- "requirements/rocm.txt"
- "requirements/build/rocm.txt"
- "requirements/test/rocm.txt"
- "setup.py"
- "csrc/"
- "cmake/"
run_all_exclude_patterns:
- "csrc/cpu/"
- "cmake/cpu_extension.cmake"
.buildkite/hardware_tests/amd.yaml
+66 -35
View File
@@ -1,42 +1,73 @@
group: Hardware - AMD Build
group: Hardware - AMD Build
steps:
- label: "AMD: :docker: build image"
key: image-build-amd
# Ensure ci_base is up-to-date before building the test image.
# Compares a content hash of ci_base-affecting files against the remote
# image label. If hashes match the build is skipped (< 30 s); if they
# differ ci_base is rebuilt and pushed automatically.
- label: "AMD: :docker: ensure ci_base"
key: ensure-ci-base-amd
depends_on: []
device: amd_cpu
no_plugin: true
commands:
- >
docker build
--build-arg max_jobs=16
--build-arg REMOTE_VLLM=1
--build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942;gfx950'
--build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
--tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
-f docker/Dockerfile.rocm
--target test
--no-cache
--progress plain .
- |
docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
command -v python3
command -v uv
command -v pytest
if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
echo No ROCm CLI found in image >&2
exit 1
fi
python3 - <<PY
import torch, vllm
print(torch.__version__)
print(vllm.__version__)
PY
echo AMD image smoke OK
'
- docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
- bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
env:
DOCKER_BUILDKIT: "1"
VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
REMOTE_VLLM: "1"
VLLM_BRANCH: "$BUILDKITE_COMMIT"
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 1
- exit_status: -10 # Agent was lost
limit: 1
- label: "AMD: :docker: build test image and artifacts"
key: image-build-amd
depends_on:
- ensure-ci-base-amd
device: amd_cpu
no_plugin: true
commands:
- |
if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
else
bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
fi
- |
docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
command -v python3
command -v uv
command -v pytest
if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
echo No ROCm CLI found in image >&2
exit 1
fi
python3 - <<PY
import torch, vllm
print(torch.__version__)
print(vllm.__version__)
PY
echo AMD image smoke OK
'
env:
DOCKER_BUILDKIT: "1"
VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"
REMOTE_VLLM: "1"
VLLM_BRANCH: "$BUILDKITE_COMMIT"
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 1
- exit_status: -10 # Agent was lost
limit: 1
.buildkite/hardware_tests/cpu.yaml
+9 -2
View File
@@ -16,6 +16,7 @@ steps:
- tests/kernels/test_onednn.py
- tests/kernels/test_awq_int4_to_int8.py
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -24,7 +25,8 @@ steps:
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
- label: CPU-Compatibility Tests
depends_on: []
@@ -62,11 +64,16 @@ steps:
source_file_dependencies:
- vllm/v1/worker/cpu/
- vllm/v1/worker/gpu/
- vllm/v1/sample/ops/topk_topp_triton.py
- vllm/v1/sample/ops/topk_topp_sampler.py
- tests/v1/sample/test_topk_topp_sampler.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
# TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
- label: CPU-Quantization Model Tests
depends_on: []
.buildkite/image_build/image_build.yaml
+14 -2
View File
@@ -6,14 +6,26 @@ steps:
timeout_in_minutes: 600
commands:
- if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
# Non-root smoke 1: the default (root) image must still be importable
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 2
- exit_status: -10 # Agent was lost
limit: 2
- label: ":docker: :smoking: Non-root smoke tests"
key: image-build-smoke-test
depends_on:
- image-build
commands:
# Smoke 1: the default (root) image must still be importable
# under a non-root UID via `--user 2000:0`. Validates the `vllm` passwd
# entry + group-0-writable /home/vllm + uv path cleanup from #31959.
# Uses `import vllm` rather than `vllm serve --help` because the latter
# instantiates `VllmConfig` which requires a GPU attached to the
# container.
- docker run --rm --user 2000:0 --entrypoint python3 "$IMAGE_TAG" -c "import vllm; print(vllm.__version__)"
# Non-root smoke 2: assert the non-root enabling invariants are baked
# Smoke 2: assert the non-root enabling invariants are baked
# into the image. Runs as UID 2000:0 via a shell so we can verify
# filesystem perms + passwd/group file state + wrapper presence without
# triggering vLLM's GPU-requiring config-init path. The opt-in
.buildkite/performance-benchmarks/tests/latency-tests-arm64-cpu.json
-1
View File
@@ -2,7 +2,6 @@
{
"test_name": "latency_llama8B_tp1",
"environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_KVCACHE_SPACE": 40
.buildkite/performance-benchmarks/tests/latency-tests-cpu.json
-1
View File
@@ -2,7 +2,6 @@
{
"test_name": "latency_llama8B_tp2",
"environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1,
.buildkite/performance-benchmarks/tests/serving-tests-arm64-cpu.json
-1
View File
@@ -13,7 +13,6 @@
200
],
"server_environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1,
.buildkite/performance-benchmarks/tests/serving-tests-cpu-asr.json
-1
View File
@@ -5,7 +5,6 @@
],
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120
},
"server_parameters": {
.buildkite/performance-benchmarks/tests/serving-tests-cpu-embed.json
-1
View File
@@ -9,7 +9,6 @@
128
],
"server_environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1,
.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json
-1
View File
@@ -5,7 +5,6 @@
],
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1,
.buildkite/performance-benchmarks/tests/serving-tests-cpu.json
-1
View File
@@ -5,7 +5,6 @@
],
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1,
.buildkite/performance-benchmarks/tests/throughput-tests-arm64-cpu.json
-1
View File
@@ -2,7 +2,6 @@
{
"test_name": "throughput_llama8B_tp1",
"environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_KVCACHE_SPACE": 40
.buildkite/performance-benchmarks/tests/throughput-tests-cpu.json
-1
View File
@@ -2,7 +2,6 @@
{
"test_name": "throughput_llama8B_tp2",
"environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
"VLLM_CPU_SGL_KERNEL": 1,
.buildkite/scripts/ci-bake-rocm.sh
+1749
View File
File diff suppressed because it is too large Load Diff
.buildkite/scripts/ci-fetch-log.sh
+21 -2
View File
@@ -9,6 +9,13 @@
# 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
@@ -26,12 +33,12 @@ if [ $# -lt 1 ]; then usage; fi
if [[ "$1" == https://* ]]; then
BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
OUT="${2:-}"
else
if [ $# -lt 2 ]; then usage; fi
BUILD="$1"
JOB="$2"
OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
OUT="${3:-}"
fi
if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
@@ -39,6 +46,18 @@ if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
usage
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)
trap 'rm -f "$COOKIES"' EXIT
.buildkite/scripts/hardware_ci/run-amd-test.sh
+121 -11
View File
@@ -52,6 +52,108 @@ cleanup_network() {
fi
}
prepare_artifact_image() {
if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
return 1
fi
if ! command -v buildkite-agent >/dev/null 2>&1; then
echo "buildkite-agent not found; cannot download ROCm wheel artifact"
return 1
fi
local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
local archive=""
local metadata_file=""
local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}"
local artifact_image=""
local artifact_key=""
local base_digest=""
local wheel_dir=""
local context_dir=""
local workspace_dir=""
artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX)
wheel_dir="${artifact_work_dir}/wheels"
context_dir="${artifact_work_dir}/context"
workspace_dir="${context_dir}/workspace"
mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}"
echo "--- Downloading ROCm wheel artifact"
if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then
echo "Failed to download ${artifact_glob}"
return 1
fi
buildkite-agent artifact download \
"artifacts/vllm-rocm-install/ci-base-image.txt" \
"${artifact_work_dir}" >/dev/null 2>&1 || true
archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1)
if [[ -z "${archive}" || ! -f "${archive}" ]]; then
echo "ROCm wheel artifact archive was not found"
return 1
fi
metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1)
if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then
base_image=$(tr -d '[:space:]' < "${metadata_file}")
fi
echo "--- Preparing local ROCm test image"
echo "Base image: ${base_image}"
docker pull "${base_image}" || return 1
base_digest=$(
docker image inspect \
--format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \
"${base_image}" 2>/dev/null || printf '%s' "${base_image}"
)
artifact_key=$(
{
printf 'base-image:%s\n' "${base_digest}"
sha256sum "${archive}"
} | sha256sum | cut -c1-24
)
artifact_image="rocm/vllm-ci-artifact:${artifact_key}"
if docker image inspect "${artifact_image}" >/dev/null 2>&1; then
echo "Using existing local ROCm artifact image: ${artifact_image}"
image_name="${artifact_image}"
return 0
fi
tar -xzf "${archive}" -C "${wheel_dir}" || return 1
if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
echo "ROCm wheel artifact did not contain a wheel"
return 1
fi
if [[ ! -d "${wheel_dir}/tests" ]]; then
echo "ROCm wheel artifact did not contain the test workspace"
return 1
fi
cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
| tar -C "${workspace_dir}" -xf - || return 1
cat > "${context_dir}/Dockerfile" <<'EOF'
ARG BASE_IMAGE
FROM ${BASE_IMAGE}
COPY wheels/ /tmp/vllm-wheels/
COPY workspace/ /vllm-workspace/
RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \
&& rm -rf /tmp/vllm-wheels
WORKDIR /vllm-workspace
EOF
echo "--- Building local ROCm test image"
docker build \
--pull=false \
--build-arg "BASE_IMAGE=${base_image}" \
-t "${artifact_image}" \
"${context_dir}" || return 1
image_name="${artifact_image}"
return 0
}
is_multi_node() {
local cmds="$1"
# Primary signal: NUM_NODES environment variable set by the pipeline
@@ -243,22 +345,30 @@ report_docker_usage
# --- Pull test image ---
echo "--- Pulling container"
image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}"
artifact_work_dir=""
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
docker pull "${image_name}"
remove_docker_container() {
# docker run uses --rm, so the container is normally already gone when the
# EXIT trap runs. Cleanup is best-effort and must not affect the test result.
docker rm -f "${container_name}" >/dev/null 2>&1 || true
if docker container inspect "${container_name}" >/dev/null 2>&1; then
docker rm -f "${container_name}" || true
fi
if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then
docker image rm -f "${image_name}" || true
else
# Keep images by default so later jobs on the same AMD node can reuse layers.
echo "Keeping ROCm test image locally: ${image_name}"
fi
if [[ -n "${artifact_work_dir}" ]]; then
rm -rf "${artifact_work_dir}"
fi
}
trap remove_docker_container EXIT
on_exit() {
local exit_code=$?
remove_docker_container
exit "$exit_code"
}
trap on_exit EXIT
if ! prepare_artifact_image; then
echo "Using full ROCm CI image: ${image_name}"
docker pull "${image_name}" || exit 1
fi
# --- Prepare commands ---
echo "--- Running container"
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+2 -1
View File
@@ -37,7 +37,8 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
.buildkite/scripts/install-kv-connectors.sh
Executable
+39
View File
@@ -0,0 +1,39 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
set -euo pipefail
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
uv pip install --system -r "${REQUIREMENTS_FILE}"
NIXL_METADATA=$(python3 - <<'PY'
import importlib.metadata as metadata
import torch
cuda_version = torch.version.cuda
if cuda_version is None:
raise SystemExit("torch.version.cuda is not set")
print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
PY
)
read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
# matching this CI image so nixl_ep_cpp links against the available libcudart.
uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
python3 - <<'PY'
import importlib.metadata as metadata
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
try:
version = metadata.version(package_name)
except metadata.PackageNotFoundError:
version = "not installed"
print(f"{package_name}: {version}")
PY
.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh
+1
View File
@@ -49,6 +49,7 @@ for BACK in "${BACKENDS[@]}"; do
--data-parallel-size 2 \
--enable-expert-parallel \
--enable-eplb \
--eplb-config '{"use_async": false}' \
--trust-remote-code \
--max-model-len 2048 \
--all2all-backend "$BACK" \
.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh
+1 -1
View File
@@ -48,7 +48,7 @@ for BACK in "${BACKENDS[@]}"; do
--enforce-eager \
--enable-eplb \
--all2all-backend "$BACK" \
--eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
--eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true, "use_async":false}' \
--tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
--data-parallel-size "${DATA_PARALLEL_SIZE}" \
--enable-expert-parallel \
.buildkite/scripts/tool_call/run-bfcl-eval.sh
+2 -2
View File
@@ -70,7 +70,7 @@ echo "============================================"
# ---- Install bfcl-eval if missing ----
if ! python3 -c "import bfcl_eval" 2>/dev/null; then
echo "Installing bfcl-eval..."
pip install "bfcl-eval>=2025.10.20.1,<2026"
uv pip install "bfcl-eval>=2025.10.20.1,<2026"
fi
# ---- Cleanup handler ----
@@ -100,7 +100,7 @@ SERVE_ARGS=(
--tensor-parallel-size "$TP_SIZE"
--max-model-len "$MAX_MODEL_LEN"
--enforce-eager
--no-enable-prefix-caching
--enable-prefix-caching
)
# Append reasoning parser if specified
.buildkite/test-amd.yaml
+20 -18
View File
@@ -1238,14 +1238,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
timeout_in_minutes: 180
@@ -1261,7 +1258,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
timeout_in_minutes: 180
@@ -1275,10 +1272,14 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
@@ -1368,7 +1369,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- label: OpenAI API correctness # TBD
timeout_in_minutes: 180
@@ -1484,7 +1485,7 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD
- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_4
@@ -1526,7 +1527,7 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_4
@@ -2745,14 +2746,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
timeout_in_minutes: 180
@@ -2768,7 +2766,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
timeout_in_minutes: 180
@@ -2782,10 +2780,14 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
@@ -2895,7 +2897,7 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_2
.buildkite/test_areas/disaggregated.yaml
+9 -9
View File
@@ -11,7 +11,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -22,7 +22,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -34,7 +34,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -46,7 +46,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
@@ -58,7 +58,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -73,7 +73,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -87,7 +87,7 @@ steps:
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
@@ -102,5 +102,5 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
.buildkite/test_areas/e2e_integration.yaml
+6 -6
View File
@@ -2,8 +2,8 @@ group: E2E Integration
depends_on:
- image-build
steps:
- label: DeepSeek V2-Lite Accuracy
key: deepseek-v2-lite-accuracy
- label: DeepSeek V2-Lite Sync EPLB Accuracy
key: deepseek-v2-lite-sync-eplb-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -12,8 +12,8 @@ steps:
commands:
- bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
- label: Qwen3-30B-A3B-FP8-block Accuracy
key: qwen3-30b-a3b-fp8-block-accuracy
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy
key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -22,8 +22,8 @@ steps:
commands:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
key: qwen3-30b-a3b-fp8-block-accuracy-b200
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200)
key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200
timeout_in_minutes: 60
device: b200-k8s
optional: true
.buildkite/test_areas/entrypoints.yaml
+9 -9
View File
@@ -11,7 +11,7 @@ steps:
- tests/entrypoints/
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
@@ -43,7 +43,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
mirror:
amd:
device: mi325_1
@@ -60,9 +60,13 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
mirror:
amd:
device: mi325_1
@@ -98,14 +102,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- 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
@@ -153,6 +154,5 @@ steps:
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
- vllm/model_executor/models/whisper.py
commands: # LMEval
- pytest -s entrypoints/openai/correctness/
.buildkite/test_areas/misc.yaml
+2 -1
View File
@@ -86,7 +86,7 @@ steps:
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# split the test to avoid interference
- pytest -v -s -m 'not cpu_test' v1/core
@@ -281,6 +281,7 @@ steps:
- vllm/model_executor/layers/quantization/quark/
- vllm/multimodal/
- vllm/outputs.py
- vllm/parser/
- vllm/platforms/
- vllm/pooling_params.py
- vllm/ray/
.buildkite/test_areas/model_executor.yaml
+9 -2
View File
@@ -14,5 +14,12 @@ steps:
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor -m '(not slow_test)'
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
# 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.
- 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
.buildkite/test_areas/model_runner_v2.yaml
+2
View File
@@ -94,11 +94,13 @@ steps:
- vllm/v1/worker/gpu_worker.py
- tests/distributed/test_pipeline_parallel.py
- tests/distributed/test_pp_cudagraph.py
- tests/v1/distributed/test_pp_dp_v2.py
commands:
- set -x
- 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
.buildkite/test_areas/rust_frontend.yaml
+6 -6
View File
@@ -16,7 +16,7 @@ steps:
- tests/benchmarks/test_serve_cli.py
- tests/entrypoints/openai/chat_completion/test_chat_completion.py
# - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
# - tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
# - tests/entrypoints/openai/completion/test_prompt_validation.py
- tests/entrypoints/openai/completion/test_shutdown.py
# - tests/entrypoints/openai/test_return_token_ids.py
@@ -28,7 +28,7 @@ steps:
- 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"
# - pytest -v -s entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
# - 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
@@ -45,19 +45,19 @@ steps:
- vllm/entrypoints/serve/
- vllm/v1/engine/
- tests/utils.py
# - tests/entrypoints/rpc/test_collective_rpc.py
# - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
- tests/entrypoints/serve/disagg/test_serving_tokens.py
- tests/entrypoints/serve/instrumentator/test_basic.py
- tests/entrypoints/serve/instrumentator/test_metrics.py
# - tests/entrypoints/serve/instrumentator/test_sleep.py
# - tests/entrypoints/serve/dev/test_sleep.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s entrypoints/rpc/test_collective_rpc.py
# - 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/instrumentator/test_sleep.py
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
- label: Rust Frontend Core Correctness
timeout_in_minutes: 30
.dockerignore
+7
View File
@@ -33,3 +33,10 @@ share/python-wheels/
*.egg
MANIFEST
rust/target/
# Not needed in Docker builds
docs/
.github/
.pre-commit-config.yaml
.clang-format
.gitattributes
format.sh
.github/CODEOWNERS
+11 -3
View File
@@ -40,6 +40,12 @@
/vllm/entrypoints/chat_utils.py @DarkLight1337
/vllm/entrypoints/llm.py @DarkLight1337
# Rust Frontend
/rust/ @BugenZhao @njhill
/build_rust.sh @BugenZhao @njhill
/rust-toolchain.toml @BugenZhao @njhill
/.buildkite/test_areas/rust* @BugenZhao @njhill
# Input/Output Processing
/vllm/sampling_params.py @njhill @NickLucche
/vllm/pooling_params.py @noooop @DarkLight1337
@@ -72,11 +78,13 @@
/vllm/v1/worker/gpu/kv_connector.py @orozery
# CI & building
/.buildkite @Harry-Chen
/docker/Dockerfile @Harry-Chen
/.buildkite @Harry-Chen @khluu
/docker/Dockerfile @Harry-Chen @khluu
/pyproject.toml @khluu
/setup.py @khluu
# Test ownership
/.buildkite/lm-eval-harness @mgoin
/.buildkite/lm-eval-harness @mgoin
/tests/distributed/test_multi_node_assignment.py @youkaichao
/tests/distributed/test_pipeline_parallel.py @youkaichao
/tests/distributed/test_same_node.py @youkaichao
.github/workflows/add_label_automerge.yml
+1 -1
View File
@@ -10,7 +10,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Add label
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
github.rest.issues.addLabels({
.github/workflows/issue_autolabel.yml
+3 -3
View File
@@ -14,7 +14,7 @@ jobs:
steps:
- name: Label issues based on keywords
id: label-step
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
// Configuration: Add new labels and keywords here
@@ -315,7 +315,7 @@ jobs:
- name: CC users for labeled issues
if: steps.label-step.outputs.labels_added != '[]'
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
// Configuration: Map labels to GitHub users to CC
@@ -392,7 +392,7 @@ jobs:
- name: Request missing ROCm info from issue author
if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const body = (context.payload.issue.body || '').toLowerCase();
.github/workflows/new_pr_bot.yml
+2 -2
View File
@@ -12,7 +12,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Update PR description
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const { owner, repo } = context.repo;
@@ -55,7 +55,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Post welcome comment for first-time contributors
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const { owner, repo } = context.repo;
.github/workflows/pre-commit.yml
+1 -1
View File
@@ -20,7 +20,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Check PR label and author merge count
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const { data: pr } = await github.rest.pulls.get({
.pre-commit-config.yaml
+1 -1
View File
@@ -21,7 +21,7 @@ repos:
rev: v21.1.2
hooks:
- id: clang-format
exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|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
.readthedocs.yaml
+1 -1
View File
@@ -9,8 +9,8 @@ build:
python: "3.12"
jobs:
post_checkout:
- bash docs/pre_run_check.sh
- git fetch origin main --unshallow --no-tags --filter=blob:none || true
- bash docs/pre_run_check.sh
pre_create_environment:
- pip install uv
create_environment:
CMakeLists.txt
+49 -25
View File
@@ -144,14 +144,14 @@ endif()
# Set up GPU language and check the torch version and warn if it isn't
# what is expected.
#
if (NOT HIP_FOUND AND CUDA_FOUND)
if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND)
set(VLLM_GPU_LANG "CUDA")
if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA})
message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} "
"expected for CUDA build, saw ${Torch_VERSION} instead.")
endif()
elseif(HIP_FOUND)
elseif(HIP_FOUND OR PYTORCH_FOUND_HIP)
set(VLLM_GPU_LANG "HIP")
# Importing torch recognizes and sets up some HIP/ROCm configuration but does
@@ -305,10 +305,6 @@ endif()
#
set(VLLM_EXT_SRC
"csrc/cache_kernels.cu"
"csrc/cache_kernels_fused.cu"
"csrc/attention/paged_attention_v1.cu"
"csrc/attention/paged_attention_v2.cu"
"csrc/cuda_view.cu"
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
"csrc/quantization/activation_kernels.cu"
@@ -637,6 +633,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/activation_kernels.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
"csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
"csrc/libtorch_stable/pos_encoding_kernels.cu"
@@ -647,7 +645,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/attention/merge_attn_states.cu"
"csrc/libtorch_stable/sampler.cu"
"csrc/libtorch_stable/topk.cu"
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
"csrc/libtorch_stable/attention/paged_attention_v1.cu"
"csrc/libtorch_stable/attention/paged_attention_v2.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"csrc/libtorch_stable/cache_kernels_fused.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_STABLE_EXT_SRC
@@ -657,8 +659,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
"csrc/libtorch_stable/permute_cols.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
set_gencode_flags_for_srcs(
@@ -683,6 +683,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"in CUDA target architectures.")
endif()
# FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0.
cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS)
set(SRCS
"csrc/libtorch_stable/fp32_router_gemm_entry.cu"
"csrc/libtorch_stable/fp32_router_gemm.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}")
else()
message(STATUS "Not building fp32_router_gemm as no compatible archs found "
"(requires SM90+ and CUDA >= 12.0).")
endif()
# Only build AllSpark kernels if we are building for at least some compatible archs.
cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
if (ALLSPARK_ARCHS)
@@ -924,13 +940,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
# nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
# from cache_kernels.cu, so it belongs in _C rather than _C_stable.
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
@@ -960,11 +974,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
@@ -1242,24 +1256,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
" in CUDA target architectures")
endif()
# DeepSeek V3 router GEMM kernel - requires SM90+
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
# DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
# (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
set(DSV3_ROUTER_GEMM_SRC
"csrc/moe/dsv3_router_gemm_entry.cu"
"csrc/moe/dsv3_router_gemm_float_out.cu"
"csrc/moe/dsv3_router_gemm_bf16_out.cu")
set_gencode_flags_for_srcs(
SRCS "${DSV3_ROUTER_GEMM_SRC}"
CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
else()
message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
" (requires SM90+ and CUDA >= 12.0)")
endif()
endif()
@@ -1286,6 +1298,14 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
"csrc/rocm/skinny_gemms.cu"
"csrc/rocm/attention.cu")
set(VLLM_ROCM_HAS_GFX1100 OFF)
if(VLLM_GPU_ARCHES MATCHES "gfx1100")
set(VLLM_ROCM_HAS_GFX1100 ON)
list(APPEND VLLM_ROCM_EXT_SRC
"csrc/rocm/q_gemm_rdna3.cu"
"csrc/rocm/q_gemm_rdna3_wmma.cu")
endif()
define_extension_target(
_rocm_C
DESTINATION vllm
@@ -1295,6 +1315,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
ARCHITECTURES ${VLLM_GPU_ARCHES}
USE_SABI 3
WITH_SOABI)
if(VLLM_ROCM_HAS_GFX1100)
target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
endif()
endif()
# Must run after the last HIP `define_extension_target` so every extension
benchmarks/kernels/benchmark_router_gemm.py
+154
View File
@@ -0,0 +1,154 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import torch.nn.functional as F
from vllm import _custom_ops as ops
from vllm.platforms import current_platform
from vllm.transformers_utils.config import get_config
from vllm.triton_utils import triton
from vllm.utils.argparse_utils import FlexibleArgumentParser
# Dimensions supported by the DSV3 specialized kernel
DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
# Dimensions supported by the gpt-oss specialized kernel
GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
# Dimensions supported by the fp32 specialized kernel (MiniMax-M2)
FP32_SUPPORTED_NUM_EXPERTS = [256]
FP32_SUPPORTED_HIDDEN_SIZES = [3072]
FP32_MAX_TOKENS = 32
def get_batch_size_range(max_batch_size):
return [2**x for x in range(14) if 2**x <= max_batch_size]
def get_model_params(config):
if config.architectures[0] in (
"DeepseekV2ForCausalLM",
"DeepseekV3ForCausalLM",
"DeepseekV32ForCausalLM",
):
num_experts = config.n_routed_experts
hidden_size = config.hidden_size
elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[
0
] in ("MiniMaxM2ForCausalLM",):
num_experts = config.num_local_experts
hidden_size = config.hidden_size
else:
raise ValueError(f"Unsupported architecture: {config.architectures}")
return num_experts, hidden_size
def get_benchmark(model, max_batch_size, trust_remote_code):
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=get_batch_size_range(max_batch_size),
x_log=False,
line_arg="provider",
line_vals=[
"torch",
"vllm",
],
line_names=["PyTorch", "vLLM"],
styles=([("blue", "-"), ("red", "-")]),
ylabel="TFLOPs",
plot_name=f"{model} router gemm throughput",
args={},
)
)
def benchmark(batch_size, provider):
config = get_config(model=model, trust_remote_code=trust_remote_code)
num_experts, hidden_size = get_model_params(config)
is_hopper_or_blackwell = current_platform.is_device_capability(
90
) or current_platform.is_device_capability_family(100)
allow_dsv3_router_gemm = (
is_hopper_or_blackwell
and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
)
allow_gpt_oss_router_gemm = (
is_hopper_or_blackwell
and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
)
is_fp32_router_model = (
is_hopper_or_blackwell
and num_experts in FP32_SUPPORTED_NUM_EXPERTS
and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES
)
allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS
# Weight dtype: fp32 kernel requires fp32 weights; others use bf16.
weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16
mat_a = torch.randn(
(batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
).contiguous()
mat_b = torch.randn(
(num_experts, hidden_size), dtype=weight_dtype, device="cuda"
).contiguous()
bias = torch.randn(
num_experts, dtype=torch.bfloat16, device="cuda"
).contiguous()
has_bias = allow_gpt_oss_router_gemm
quantiles = [0.5, 0.2, 0.8]
if provider == "torch":
def runner():
if allow_fp32_router_gemm:
F.linear(mat_a.float(), mat_b)
elif has_bias:
F.linear(mat_a, mat_b, bias)
else:
F.linear(mat_a, mat_b)
elif provider == "vllm":
def runner():
if allow_dsv3_router_gemm:
ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
elif allow_fp32_router_gemm:
ops.fp32_router_gemm(mat_a, mat_b)
elif allow_gpt_oss_router_gemm:
ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
elif is_fp32_router_model:
# batch_size > FP32_MAX_TOKENS: fall back to F.linear
F.linear(mat_a.float(), mat_b)
else:
F.linear(mat_a, mat_b)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
runner, quantiles=quantiles
)
def tflops(t_ms):
flops = 2 * batch_size * hidden_size * num_experts
return flops / (t_ms * 1e-3) / 1e12
return tflops(ms), tflops(max_ms), tflops(min_ms)
return benchmark
if __name__ == "__main__":
parser = FlexibleArgumentParser()
parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
parser.add_argument("--max-batch-size", default=16, type=int)
parser.add_argument("--trust-remote-code", action="store_true")
args = parser.parse_args()
# Get the benchmark function
benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
# Run performance benchmark
benchmark.run(print_data=True)
cmake/cpu_extension.cmake
+31 -1
View File
@@ -369,6 +369,18 @@ else()
add_compile_definitions(-DVLLM_NUMA_DISABLED)
endif()
# check if the pytorch wheel ships libopenblas.so.
set(VLLM_OPENBLAS_LIB "")
if (NOT ENABLE_X86_ISA)
file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
"${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
# Note: we don't link openblas directly to _C extension, as it's available through libtorch.so
if (_VLLM_TORCH_OPENBLAS_LIBS)
list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
endif()
endif()
#
# Generate CPU attention dispatch header
#
@@ -387,6 +399,7 @@ endif()
#
set(VLLM_EXT_SRC
"csrc/cpu/activation.cpp"
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/layernorm.cpp"
@@ -396,6 +409,13 @@ set(VLLM_EXT_SRC
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp")
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
set(VLLM_EXT_SRC
"csrc/cpu/cpu_wna16.cpp"
${VLLM_EXT_SRC})
endif()
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
@@ -403,6 +423,12 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
${VLLM_EXT_SRC})
endif()
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
${VLLM_EXT_SRC})
endif()
if(USE_ONEDNN)
set(VLLM_EXT_SRC
"csrc/cpu/dnnl_kernels.cpp"
@@ -411,7 +437,6 @@ endif()
if (ENABLE_X86_ISA)
set(VLLM_EXT_SRC_SGL
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/sgl-kernels/conv.cpp"
"csrc/cpu/sgl-kernels/gemm.cpp"
"csrc/cpu/sgl-kernels/gemm_int8.cpp"
@@ -423,6 +448,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/sgl-kernels/moe_fp8.cpp")
set(VLLM_EXT_SRC_AVX512
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/shm.cpp"
"csrc/cpu/cpu_wna16.cpp"
"csrc/cpu/cpu_fused_moe.cpp"
@@ -439,6 +465,7 @@ if (ENABLE_X86_ISA)
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
set(VLLM_EXT_SRC_AVX2
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
@@ -512,6 +539,9 @@ else()
USE_SABI 3
WITH_SOABI
)
if (VLLM_OPENBLAS_LIB)
target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
endif()
endif()
message(STATUS "Enabling C extension.")
cmake/external_projects/vllm_flash_attn.cmake
+2 -2
View File
@@ -31,7 +31,7 @@ endif()
if(VLLM_FLASH_ATTN_SRC_DIR)
FetchContent_Declare(
vllm-flash-attn SOURCE_DIR
vllm-flash-attn SOURCE_DIR
${VLLM_FLASH_ATTN_SRC_DIR}
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
)
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
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
cmake/hipify.py
+31 -10
View File
@@ -14,7 +14,18 @@ import argparse
import os
import shutil
from torch.utils.hipify.hipify_python import hipify
from torch.utils.hipify.hipify_python import get_hip_file_path, hipify
def _expected_hip_build_path(source_abs: str, output_directory: str) -> str:
"""Match torch.utils.hipify.hipify_python.preprocessor fout_path naming."""
rel = os.path.relpath(source_abs, output_directory)
return os.path.abspath(
os.path.join(
output_directory, get_hip_file_path(rel, is_pytorch_extension=True)
)
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
@@ -53,7 +64,11 @@ if __name__ == "__main__":
hipify_result = hipify(
project_directory=args.project_dir,
output_directory=args.output_dir,
header_include_dirs=[],
# Hipify resolves quoted includes next to the including file first; vLLM
# uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h"
# from quantization/w8a8/fp8/*.cu). Without an include root here, those
# headers are never found and are not hipified or rewritten in dependents.
header_include_dirs=["."],
includes=includes,
extra_files=extra_files,
show_detailed=True,
@@ -64,14 +79,20 @@ if __name__ == "__main__":
hipified_sources = []
for source in args.sources:
s_abs = os.path.abspath(source)
hipified_s_abs = (
hipify_result[s_abs].hipified_path
if (
s_abs in hipify_result
and hipify_result[s_abs].hipified_path is not None
)
else s_abs
)
if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:
path = hipify_result[s_abs].hipified_path
# PyTorch skips writing when is_pytorch_extension and text unchanged;
# hipified_path then stays *.cu. CMake expects *.hip under output_dir.
if s_abs.endswith(".cu") and path.endswith(".cu"):
dest = _expected_hip_build_path(s_abs, args.output_dir)
if os.path.normpath(path) != os.path.normpath(dest):
os.makedirs(os.path.dirname(dest), exist_ok=True)
shutil.copy2(path, dest)
hipified_s_abs = dest
else:
hipified_s_abs = path
else:
hipified_s_abs = s_abs
hipified_sources.append(hipified_s_abs)
assert len(hipified_sources) == len(args.sources)
cmake/utils.cmake
+18
View File
@@ -81,6 +81,14 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})
# Chain hipify targets so they run sequentially. Parallel hipify
# invocations race on shutil.copytree, overwriting .hip files
# produced by another target back to .cu originals.
if (DEFINED _VLLM_LAST_HIPIFY_TARGET)
add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET})
endif()
set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE)
# Swap out original extension sources with hipified sources.
list(APPEND HIP_SRCS ${CXX_SRCS})
set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
@@ -476,6 +484,16 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
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" "${TGT_CUDA_ARCHS}")
else()
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()
#
# Override the GPU architectures detected by cmake/torch and filter them by
# `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in
csrc/cpu/cpu_fused_moe.cpp
+49 -1
View File
@@ -30,7 +30,12 @@
}()
namespace {
enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul };
enum class FusedMOEAct {
SiluAndMul,
SwigluOAIAndMul,
GeluAndMul,
GeluTanhAndMul,
};
FusedMOEAct get_act_type(const std::string& act) {
if (act == "silu") {
@@ -39,6 +44,8 @@ FusedMOEAct get_act_type(const std::string& act) {
return FusedMOEAct::SwigluOAIAndMul;
} else if (act == "gelu") {
return FusedMOEAct::GeluAndMul;
} else if (act == "gelu_tanh") {
return FusedMOEAct::GeluTanhAndMul;
} else {
TORCH_CHECK(false, "Invalid act type: " + act);
}
@@ -143,6 +150,44 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
}
}
template <typename scalar_t>
void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride,
const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
vec_op::FP32Vec16 w2_vec(0.5);
vec_op::FP32Vec16 w3_vec(0.044715);
alignas(64) float temp[16];
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
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);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
float* __restrict__ input,
@@ -160,6 +205,9 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
case FusedMOEAct::GeluAndMul:
gelu_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::GeluTanhAndMul:
gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
return;
default:
TORCH_CHECK(false, "Unsupported act type.");
}
csrc/cpu/cpu_types_riscv_impl.hpp
+51
View File
@@ -94,6 +94,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
explicit FP16Vec16(const c10::Half v)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
explicit FP16Vec16(const FP32Vec16& vec);
void save(void* ptr) const {
@@ -165,6 +169,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
explicit BF16Vec16(const c10::BFloat16 v)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
@@ -290,6 +297,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
}
explicit BF16Vec16(const c10::BFloat16 v)
: reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
VEC_ELEM_NUM)) {}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
float tmp[16];
@@ -629,6 +639,19 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
data.reg, data.reg)) {};
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
const uint64_t q_values = static_cast<uint64_t>(value);
auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
auto shifts =
RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
auto shifted =
RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
auto idx64 =
RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
}
explicit FP32Vec16(const FP16Vec16& v);
#ifdef __riscv_zvfbfmin
@@ -641,6 +664,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
#endif
// FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
// load_b_pair_vec template to compile on all platforms.
explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
FP32Vec16 operator+(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -891,6 +918,30 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
acc = acc.fma(a, b);
}
template <typename VecT>
static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
vec0.save(values0);
vec1.save(values1);
auto* packed = reinterpret_cast<uint32_t*>(ptr);
for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
packed[i] = static_cast<uint32_t>(values0[i]) |
(static_cast<uint32_t>(values1[i]) << 16);
}
}
static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
void* ptr) {
interleave_save_16b(vec0, vec1, ptr);
}
static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
void* ptr) {
interleave_save_16b(vec0, vec1, ptr);
}
#ifdef __riscv_zvfbfmin
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
csrc/cpu/cpu_types_vsx.hpp
+106 -1
View File
@@ -89,6 +89,35 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
}
};
struct FP16Vec16 : public Vec<FP16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
ss16x8x2_t reg;
explicit FP16Vec16(const void* ptr) {
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
}
explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
explicit FP16Vec16(const FP32Vec16&);
void save(void* ptr) const {
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst(reg.val[1], 16, (signed short*)ptr);
}
void save(void* ptr, int elem_num) const {
int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
if (num <= 8) {
vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
} else {
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
}
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
@@ -100,6 +129,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
}
explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
@@ -379,6 +410,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = vec_xl(48, ptr);
}
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
explicit FP32Vec16(const FP32Vec16& data) {
@@ -402,6 +435,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = data.reg.val[1];
}
explicit FP32Vec16(const FP16Vec16& v);
explicit FP32Vec16(const BF16Vec16& v) {
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -735,6 +769,40 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
#endif
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
alignas(16) float temp_fp32[16];
alignas(16) c10::Half temp_fp16[16];
vec_xst(v.reg.val[0], 0, temp_fp32);
vec_xst(v.reg.val[1], 16, temp_fp32);
vec_xst(v.reg.val[2], 32, temp_fp32);
vec_xst(v.reg.val[3], 48, temp_fp32);
for (int i = 0; i < 16; i++) {
temp_fp16[i] = c10::Half(temp_fp32[i]);
}
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
}
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
alignas(16) c10::Half temp_fp16[16];
alignas(16) float temp_fp32[16];
vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
for (int i = 0; i < 16; i++) {
temp_fp32[i] = float(temp_fp16[i]);
}
reg.val[0] = vec_xl(0, temp_fp32);
reg.val[1] = vec_xl(16, temp_fp32);
reg.val[2] = vec_xl(32, temp_fp32);
reg.val[3] = vec_xl(48, temp_fp32);
}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
#ifdef _ARCH_PWR10
__vector signed short ret[4];
@@ -794,6 +862,43 @@ inline void prefetch(const void* addr) {
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
}
}; // namespace vec_op
struct INT8Vec64 {
__vector signed char data[4];
INT8Vec64() = default;
explicit INT8Vec64(const int8_t* ptr) {
data[0] = vec_xl(0, ptr);
data[1] = vec_xl(16, ptr);
data[2] = vec_xl(32, ptr);
data[3] = vec_xl(48, ptr);
}
explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
void save(int8_t* ptr) const {
vec_xst(data[0], 0, ptr);
vec_xst(data[1], 16, ptr);
vec_xst(data[2], 32, ptr);
vec_xst(data[3], 48, ptr);
}
void save(int8_t* ptr, int elem_num) const {
if (elem_num <= 0) return;
int full_vecs = elem_num / 16;
for (int i = 0; i < full_vecs && i < 4; i++) {
vec_xst(data[i], i * 16, ptr);
}
int remaining = elem_num % 16;
if (remaining > 0 && full_vecs < 4) {
vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
}
}
void nt_save(int8_t* ptr) const { save(ptr); }
};
} // namespace vec_op
#endif
csrc/cpu/sgl-kernels/blas_gemm.h
+82
View File
@@ -0,0 +1,82 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#include <ATen/native/CPUBlas.h>
// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
#if defined(VLLM_HAS_OPENBLAS)
extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
float* alpha, const at::BFloat16* a, int* lda,
const at::BFloat16* b, int* ldb, float* beta, float* c,
int* ldc);
extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
float* alpha, const float* a, int* lda, const float* b,
int* ldb, float* beta, float* c, int* ldc);
inline char blas_transpose(at::native::TransposeType trans) {
switch (trans) {
case at::native::TransposeType::NoTranspose:
return 'n';
case at::native::TransposeType::Transpose:
return 't';
case at::native::TransposeType::ConjTranspose:
return 'c';
}
return 'n';
}
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const at::BFloat16* a,
int64_t lda, const at::BFloat16* b, int64_t ldb,
float beta, float* c, int64_t ldc) {
char transa_ = blas_transpose(transa);
char transb_ = blas_transpose(transb);
int m_ = static_cast<int>(m);
int n_ = static_cast<int>(n);
int k_ = static_cast<int>(k);
int lda_ = static_cast<int>(lda);
int ldb_ = static_cast<int>(ldb);
int ldc_ = static_cast<int>(ldc);
sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
c, &ldc_);
}
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const float* a, int64_t lda,
const float* b, int64_t ldb, float beta, float* c,
int64_t ldc) {
char transa_ = blas_transpose(transa);
char transb_ = blas_transpose(transb);
int m_ = static_cast<int>(m);
int n_ = static_cast<int>(n);
int k_ = static_cast<int>(k);
int lda_ = static_cast<int>(lda);
int ldb_ = static_cast<int>(ldb);
int ldc_ = static_cast<int>(ldc);
sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
c, &ldc_);
}
inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
int64_t, int64_t, int64_t, float, const at::Half*,
int64_t, const at::Half*, int64_t, float, float*,
int64_t) {
TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
}
#else
template <typename scalar_t>
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const scalar_t* a, int64_t lda,
const scalar_t* b, int64_t ldb, float beta, float* c,
int64_t ldc) {
auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
}
#endif
csrc/cpu/sgl-kernels/fla.cpp
+278 -141
View File
@@ -301,25 +301,42 @@ void chunk_gated_delta_rule_kernel_impl(
// attn = k_beta @ key.transpose(-1, -2)
// attn: [B, HV, num_chunk, chunk_size, chunk_size]
// transpose and pack for key
pack_vnni<scalar_t>(
/* dst */ k_transpose,
/* src */ curr_k_pad,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// k_beta @ key.transpose(-1, -2)
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ curr_k_beta,
/* B */ k_transpose,
/* C */ curr_attn);
if constexpr (brgemm_supported()) {
pack_vnni<scalar_t>(
/* dst */ k_transpose,
/* src */ curr_k_pad,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// k_beta @ key.transpose(-1, -2)
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ curr_k_beta,
/* B */ k_transpose,
/* C */ curr_attn);
} else {
blas_gemm(
at::native::TransposeType::Transpose,
at::native::TransposeType::NoTranspose,
chunk_size,
chunk_size,
qk_head_size,
1.0f,
curr_k_pad,
qk_head_size,
curr_k_beta,
qk_head_size,
0.0f,
curr_attn,
chunk_size);
}
// attn = attn * decay_mask
for (int64_t m = 0; m < chunk_size; m++) {
at::vec::map2<float>(
@@ -413,25 +430,42 @@ void chunk_gated_delta_rule_kernel_impl(
// k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
// k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
// pack for value
pack_vnni2<scalar_t>(
/* dst */ v_pack,
/* src */ curr_v_beta,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// value = attn @ v_beta
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ v_pack,
/* C */ curr_value);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ v_pack,
/* src */ curr_v_beta,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// value = attn @ v_beta
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ v_pack,
/* C */ curr_value);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
chunk_size,
1.0f,
curr_v_beta,
v_head_size,
curr_attn_reduced,
chunk_size,
0.0f,
curr_value,
v_head_size);
}
// k_beta_g = k_beta * g.exp().unsqueeze(-1)
for (int64_t j = 0; j < chunk_size; j++) {
int64_t i = 0;
@@ -445,25 +479,42 @@ void chunk_gated_delta_rule_kernel_impl(
}
}
// pack for k_beta_g
pack_vnni2<scalar_t>(
/* dst */ k_beta_g_pack,
/* src */ k_beta_g,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ qk_head_size);
// k_cumdecay = attn @ k_beta_g
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ qk_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ qk_head_size,
/* ldc */ qk_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ k_beta_g_pack,
/* C */ k_cumdecay);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ k_beta_g_pack,
/* src */ k_beta_g,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ qk_head_size);
// k_cumdecay = attn @ k_beta_g
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ qk_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ qk_head_size,
/* ldc */ qk_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ k_beta_g_pack,
/* C */ k_cumdecay);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
qk_head_size,
chunk_size,
chunk_size,
1.0f,
k_beta_g,
qk_head_size,
curr_attn_reduced,
chunk_size,
0.0f,
k_cumdecay,
qk_head_size);
}
for (int i = 0; i < chunk_size; i++) {
at::vec::map<scalar_t>(
[](fVec x) { return x; },
@@ -551,25 +602,42 @@ void chunk_gated_delta_rule_kernel_impl(
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
// k_transpose_i = k_i.transpose(-1, -2)
pack_vnni<scalar_t>(
/* dst */ k_transpose_i,
/* src */ k_i,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// attn_i = q_i @ k_transpose_i
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ q_i,
/* B */ k_transpose_i,
/* C */ attn_i);
if constexpr (brgemm_supported()) {
pack_vnni<scalar_t>(
/* dst */ k_transpose_i,
/* src */ k_i,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// attn_i = q_i @ k_transpose_i
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ q_i,
/* B */ k_transpose_i,
/* C */ attn_i);
} else {
blas_gemm(
at::native::TransposeType::Transpose,
at::native::TransposeType::NoTranspose,
chunk_size,
chunk_size,
qk_head_size,
1.0f,
k_i,
qk_head_size,
q_i,
qk_head_size,
0.0f,
attn_i,
chunk_size);
}
// attn_i = attn_i * decay_mask_i
for (int64_t m = 0; m < chunk_size; m++) {
auto attn_i_m = attn_i + m * chunk_size;
@@ -609,28 +677,45 @@ void chunk_gated_delta_rule_kernel_impl(
}
// pack for curr_last_recurrent_state
pack_vnni2<scalar_t>(
/* dst */ curr_last_recurrent_state_pack_reduced,
/* src */ curr_last_recurrent_state_reduced,
/* N */ qk_head_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ curr_last_recurrent_state_pack_reduced,
/* src */ curr_last_recurrent_state_reduced,
/* N */ qk_head_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
// k_cumdecay_i: [chunk_size, EK]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ k_cumdecay_i_reduced,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ v_prime);
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
// k_cumdecay_i: [chunk_size, EK]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ k_cumdecay_i_reduced,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ v_prime);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
qk_head_size,
1.0f,
curr_last_recurrent_state_reduced,
v_head_size,
k_cumdecay_i_reduced,
qk_head_size,
0.0f,
v_prime,
v_head_size);
}
// v_new = v_prime = v_i - v_prime
// v_i: [chunk_size, EV]
@@ -663,41 +748,75 @@ void chunk_gated_delta_rule_kernel_impl(
}
// attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ qg,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ attn_inter);
if constexpr (brgemm_supported()) {
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ qg,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ attn_inter);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
qk_head_size,
1.0f,
curr_last_recurrent_state_reduced,
v_head_size,
qg,
qk_head_size,
0.0f,
attn_inter,
v_head_size);
}
// core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
// pack for v_prime
pack_vnni2<scalar_t>(
/* dst */ v_prime_pack_reduced,
/* src */ v_prime_reduced,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
// attn_i: [chunk_size, chunk_size]
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ true,
/* A */ attn_i_reduced,
/* B */ v_prime_pack_reduced,
/* C */ attn_inter);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ v_prime_pack_reduced,
/* src */ v_prime_reduced,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
// attn_i: [chunk_size, chunk_size]
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ true,
/* A */ attn_i_reduced,
/* B */ v_prime_pack_reduced,
/* C */ attn_inter);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
chunk_size,
1.0f,
v_prime_reduced,
v_head_size,
attn_i_reduced,
chunk_size,
1.0f,
attn_inter,
v_head_size);
}
// core_attn_out[:, :, i] = attn_inter
for (int64_t m = 0; m < chunk_size; m++) {
@@ -762,17 +881,34 @@ void chunk_gated_delta_rule_kernel_impl(
/* ld_dst */ chunk_size);
// kgv = kg.transpose(-1, -2) @ v_new
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ qk_head_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ kg_transpose,
/* B */ v_prime_pack_reduced,
/* C */ kgv);
if constexpr (brgemm_supported()) {
at::native::cpublas::brgemm(
/* M */ qk_head_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ kg_transpose,
/* B */ v_prime_pack_reduced,
/* C */ kgv);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
qk_head_size,
chunk_size,
1.0f,
v_prime_reduced,
v_head_size,
kg_transpose,
chunk_size,
0.0f,
kgv,
v_head_size);
}
// last_recurrent_state = 1) + 2)
for (int64_t m = 0; m < qk_head_size; m++) {
at::vec::map2<float>(
@@ -921,7 +1057,8 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
// See: https://github.com/sgl-project/sglang/pull/26634
float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
fVec beta_vec = fVec(beta_val);
int64_t dvi = 0;
for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
csrc/cpu/sgl-kernels/gemm.h
+18 -7
View File
@@ -4,9 +4,12 @@
// clang-format off
#pragma once
#include <ATen/native/CPUBlas.h>
#include "common.h"
#include "blas_gemm.h"
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
#define CPU_CAPABILITY_AVX512
#endif
// amx-bf16
#define TILE_M 16
@@ -21,31 +24,39 @@ constexpr int block_size_n() {
return 2 * TILE_N;
}
constexpr bool brgemm_supported() {
#if defined(CPU_CAPABILITY_AVX512)
return true;
#else
return false;
#endif
}
// define threshold using brgemm (intel AMX)
template <typename T>
inline bool can_use_brgemm(int M);
template <>
inline bool can_use_brgemm<at::BFloat16>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
template <>
inline bool can_use_brgemm<at::Half>(int M) {
return true;
return brgemm_supported();
}
// this requires PyTorch 2.7 or above
template <>
inline bool can_use_brgemm<int8_t>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
template <>
inline bool can_use_brgemm<uint8_t>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
template <>
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
// work around compiler internal error
csrc/cpu/sgl-kernels/vec.h
+2
View File
@@ -11,7 +11,9 @@
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
#if defined(CPU_CAPABILITY_AVX512)
#include <immintrin.h>
#endif
namespace {
using namespace at::vec;
csrc/cpu/shm.cpp
+10 -8
View File
@@ -5,7 +5,7 @@
#include <sys/stat.h>
#include <unistd.h>
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
#include <atomic>
#endif
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
};
struct ThreadSHMContext {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
// memory model is weaker on AArch64, so we use atomic variables for
// consumer (load-acquire) and producer (store-release) to make sure
// that a stamp cannot be ready before the corresponding data is ready.
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
TORCH_CHECK((size_t)this % 64 == 0);
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
}
char get_curr_stamp(int idx) const {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
#else
return _curr_thread_stamp[idx];
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
}
char get_ready_stamp(int idx) const {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
#else
return _ready_thread_stamp[idx];
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
}
void next_stamp() {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
1, std::memory_order_release);
#else
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
}
void commit_ready_stamp() {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
_ready_thread_stamp[local_stamp_buffer_idx].store(
_curr_thread_stamp[local_stamp_buffer_idx].load(
std::memory_order_relaxed),
@@ -186,8 +186,10 @@ struct ThreadSHMContext {
break;
}
++_spinning_count;
#ifdef __aarch64__
#if defined(__aarch64__)
__asm__ __volatile__("yield");
#elif defined(__powerpc64__)
__asm__ __volatile__("or 1,1,1");
#else
_mm_pause();
#endif // __aarch64__
csrc/cpu/torch_bindings.cpp
+22 -21
View File
@@ -378,7 +378,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#endif
// SHM CCL
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
defined(__powerpc64__)
ops.def(
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
"int",
@@ -447,6 +448,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"bool is_vnni) -> Tensor");
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
// Adapted from sglang: casual_conv1d kernels
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
&causal_conv1d_weight_pack);
ops.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
"conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
"int pad_slot_id, bool is_vnni) -> "
"Tensor");
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// Adapted from sglang: GDN kernels
ops.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -470,25 +490,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"-> (Tensor, Tensor)");
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
// Adapted from sglang: casual_conv1d kernels
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
&causal_conv1d_weight_pack);
ops.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
"conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
"int pad_slot_id, bool is_vnni) -> "
"Tensor");
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// CPU attention kernels
ops.def(
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
@@ -518,7 +519,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
// WNA16
#if defined(__AVX512F__)
#if defined(__AVX512F__) || defined(__riscv_v)
ops.def(
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
csrc/libtorch_stable/activation_kernels.cu
+1 -1
View File
@@ -4,7 +4,7 @@
#include <cmath>
#include "../cuda_compat.h"
#include "../cuda_vec_utils.cuh"
#include "cuda_vec_utils.cuh"
#include "dispatch_utils.h"
#include "torch_utils.h"
csrc/async_util.cuh → csrc/libtorch_stable/async_util.cuh
View File
csrc/attention/attention_kernels.cuh → csrc/libtorch_stable/attention/attention_kernels.cuh
+4 -7
View File
@@ -17,21 +17,18 @@
* limitations under the License.
*/
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <algorithm>
#include "attention_dtypes.h"
#include "../../attention/attention_dtypes.h"
#include "attention_utils.cuh"
#include "../cuda_compat.h"
#include "../../cuda_compat.h"
#ifdef USE_ROCM
#include <hip/hip_bf16.h>
#include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
#include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
typedef __hip_bfloat16 __nv_bfloat16;
#else
#include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
#include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
#endif
#define MAX(a, b) ((a) > (b) ? (a) : (b))
csrc/attention/attention_utils.cuh → csrc/libtorch_stable/attention/attention_utils.cuh
+2 -2
View File
@@ -18,8 +18,8 @@
*/
#pragma once
#include "../cuda_compat.h"
#include "attention_dtypes.h"
#include "../../cuda_compat.h"
#include "../../attention/attention_dtypes.h"
#include <float.h>
#include <type_traits>
csrc/libtorch_stable/attention/merge_attn_states.cu
+1 -1
View File
@@ -7,7 +7,7 @@
#include <torch/headeronly/core/ScalarType.h>
#include "../../attention/attention_dtypes.h"
#include "../../attention/attention_utils.cuh"
#include "attention_utils.cuh"
#include "../../quantization/w8a8/fp8/common.cuh"
namespace vllm {
csrc/attention/paged_attention_v1.cu → csrc/libtorch_stable/attention/paged_attention_v1.cu
+41 -37
View File
@@ -16,8 +16,9 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "../torch_utils.h"
#include "attention_kernels.cuh"
#include "../cuda_compat.h"
#include "../../cuda_compat.h"
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -44,13 +45,15 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
int NUM_THREADS = 128>
void paged_attention_v1_launcher(
torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int num_kv_heads, float scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
torch::Tensor& v_scale, const int tp_rank,
const int blocksparse_local_blocks, const int blocksparse_vert_stride,
const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
torch::stable::Tensor& out, torch::stable::Tensor& query,
torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
torch::stable::Tensor& seq_lens, int max_seq_len,
const std::optional<torch::stable::Tensor>& alibi_slopes,
torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
const int tp_rank, const int blocksparse_local_blocks,
const int blocksparse_vert_stride, const int blocksparse_block_size,
const int blocksparse_head_sliding_step) {
int num_seqs = query.size(0);
int num_heads = query.size(1);
int head_size = query.size(2);
@@ -69,8 +72,8 @@ void paged_attention_v1_launcher(
T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
int* block_tables_ptr = block_tables.data_ptr<int>();
int* seq_lens_ptr = seq_lens.data_ptr<int>();
int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
@@ -85,8 +88,9 @@ void paged_attention_v1_launcher(
dim3 grid(num_heads, num_seqs, 1);
dim3 block(NUM_THREADS);
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const torch::stable::accelerator::DeviceGuard device_guard(
query.get_device_index());
const cudaStream_t stream = get_current_cuda_stream();
switch (head_size) {
// NOTE(woosuk): To reduce the compilation time, we only compile for the
// head sizes that we use in the model. However, we can easily extend this
@@ -119,7 +123,7 @@ void paged_attention_v1_launcher(
LAUNCH_PAGED_ATTENTION_V1(256);
break;
default:
TORCH_CHECK(false, "Unsupported head size: ", head_size);
STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
break;
}
}
@@ -141,43 +145,43 @@ void paged_attention_v1_launcher(
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
switch (block_size) { \
case 8: \
CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
break; \
case 16: \
CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
break; \
case 32: \
CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
break; \
default: \
TORCH_CHECK(false, "Unsupported block size: ", block_size); \
break; \
#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
switch (block_size) { \
case 8: \
CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
break; \
case 16: \
CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
break; \
case 32: \
CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
break; \
default: \
STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
break; \
}
void paged_attention_v1(
torch::Tensor& out, // [num_seqs, num_heads, head_size]
torch::Tensor& query, // [num_seqs, num_heads, head_size]
torch::Tensor&
torch::stable::Tensor& out, // [num_seqs, num_heads, head_size]
torch::stable::Tensor& query, // [num_seqs, num_heads, head_size]
torch::stable::Tensor&
key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
torch::Tensor&
torch::stable::Tensor&
value_cache, // [num_blocks, num_heads, head_size, block_size]
int64_t num_kv_heads, // [num_heads]
double scale,
torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
torch::Tensor& seq_lens, // [num_seqs]
torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
torch::stable::Tensor& seq_lens, // [num_seqs]
int64_t block_size, int64_t max_seq_len,
const std::optional<torch::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale, const int64_t tp_rank,
const std::optional<torch::stable::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale, const int64_t tp_rank,
const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step) {
const bool is_block_sparse = (blocksparse_vert_stride > 1);
DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
CALL_V1_LAUNCHER_BLOCK_SIZE)
}
csrc/attention/paged_attention_v2.cu → csrc/libtorch_stable/attention/paged_attention_v2.cu
+47 -41
View File
@@ -16,8 +16,9 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "../torch_utils.h"
#include "attention_kernels.cuh"
#include "../cuda_compat.h"
#include "../../cuda_compat.h"
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -44,14 +45,16 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
int NUM_THREADS = 128, int PARTITION_SIZE = 512>
void paged_attention_v2_launcher(
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int num_kv_heads, float scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
torch::Tensor& v_scale, const int tp_rank,
const int blocksparse_local_blocks, const int blocksparse_vert_stride,
const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
const int tp_rank, const int blocksparse_local_blocks,
const int blocksparse_vert_stride, const int blocksparse_block_size,
const int blocksparse_head_sliding_step) {
int num_seqs = query.size(0);
int num_heads = query.size(1);
int head_size = query.size(2);
@@ -73,8 +76,8 @@ void paged_attention_v2_launcher(
T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
int* block_tables_ptr = block_tables.data_ptr<int>();
int* seq_lens_ptr = seq_lens.data_ptr<int>();
int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
@@ -91,8 +94,9 @@ void paged_attention_v2_launcher(
int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
dim3 block(NUM_THREADS);
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const torch::stable::accelerator::DeviceGuard device_guard(
query.get_device_index());
const cudaStream_t stream = get_current_cuda_stream();
switch (head_size) {
// NOTE(woosuk): To reduce the compilation time, we only compile for the
// head sizes that we use in the model. However, we can easily extend this
@@ -125,7 +129,7 @@ void paged_attention_v2_launcher(
LAUNCH_PAGED_ATTENTION_V2(256);
break;
default:
TORCH_CHECK(false, "Unsupported head size: ", head_size);
STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
break;
}
}
@@ -148,46 +152,48 @@ void paged_attention_v2_launcher(
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
switch (block_size) { \
case 8: \
CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
break; \
case 16: \
CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
break; \
case 32: \
CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
break; \
default: \
TORCH_CHECK(false, "Unsupported block size: ", block_size); \
break; \
#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
switch (block_size) { \
case 8: \
CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
break; \
case 16: \
CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
break; \
case 32: \
CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
break; \
default: \
STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
break; \
}
void paged_attention_v2(
torch::Tensor& out, // [num_seqs, num_heads, head_size]
torch::Tensor& exp_sums, // [num_seqs, num_heads, max_num_partitions]
torch::Tensor& max_logits, // [num_seqs, num_heads, max_num_partitions]
torch::Tensor&
torch::stable::Tensor& out, // [num_seqs, num_heads, head_size]
torch::stable::Tensor&
exp_sums, // [num_seqs, num_heads, max_num_partitions]
torch::stable::Tensor&
max_logits, // [num_seqs, num_heads, max_num_partitions]
torch::stable::Tensor&
tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
torch::Tensor& query, // [num_seqs, num_heads, head_size]
torch::Tensor&
torch::stable::Tensor& query, // [num_seqs, num_heads, head_size]
torch::stable::Tensor&
key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
torch::Tensor&
torch::stable::Tensor&
value_cache, // [num_blocks, num_heads, head_size, block_size]
int64_t num_kv_heads, // [num_heads]
double scale,
torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
torch::Tensor& seq_lens, // [num_seqs]
torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
torch::stable::Tensor& seq_lens, // [num_seqs]
int64_t block_size, int64_t max_seq_len,
const std::optional<torch::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale, const int64_t tp_rank,
const std::optional<torch::stable::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale, const int64_t tp_rank,
const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step) {
const bool is_block_sparse = (blocksparse_vert_stride > 1);
DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
CALL_V2_LAUNCHER_BLOCK_SIZE)
}
csrc/cache_kernels.cu → csrc/libtorch_stable/cache_kernels.cu
+358 -309
View File
File diff suppressed because it is too large Load Diff
csrc/cache_kernels_fused.cu → csrc/libtorch_stable/cache_kernels_fused.cu
+102 -89
View File
@@ -1,15 +1,13 @@
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "cuda_compat.h"
#include "torch_utils.h"
#include "dispatch_utils.h"
#include "quantization/w8a8/fp8/common.cuh"
#include "../cuda_compat.h"
#include "../quantization/w8a8/fp8/common.cuh"
#ifdef USE_ROCM
#include "quantization/w8a8/fp8/amd/quant_utils.cuh"
#include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
#else
#include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
#include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
#endif
#ifdef USE_ROCM
@@ -164,43 +162,52 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
} // namespace vllm
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
do { \
VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
using qk_t = scalar_t; \
VLLM_DISPATCH_FLOATING_TYPES( \
rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
[&] { \
using cos_sin_t = scalar_t; \
if (rope_is_neox) { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} else { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} \
}); \
}); \
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
do { \
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
q_pe.scalar_type(), "qk_scalar_type", [&] { \
using qk_t = scalar_t; \
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
rope_cos_sin_cache.scalar_type(), \
"rope_cos_sin_cache_scalar_type", [&] { \
using cos_sin_t = scalar_t; \
if (rope_is_neox) { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.const_data_ptr<int64_t>(), \
q_pe.mutable_data_ptr<qk_t>(), \
k_pe.mutable_data_ptr<qk_t>(), \
kv_c.const_data_ptr<qk_t>(), \
rope_cos_sin_cache.const_data_ptr<cos_sin_t>(), \
rot_dim, q_pe_stride_token, q_pe_stride_head, \
k_pe_stride, kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>( \
kv_cache.mutable_data_ptr()), \
slot_mapping.const_data_ptr<int64_t>(), \
block_stride, entry_stride, kv_lora_rank, \
block_size, \
kv_cache_quant_scale.const_data_ptr<float>()); \
} else { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.const_data_ptr<int64_t>(), \
q_pe.mutable_data_ptr<qk_t>(), \
k_pe.mutable_data_ptr<qk_t>(), \
kv_c.const_data_ptr<qk_t>(), \
rope_cos_sin_cache.const_data_ptr<cos_sin_t>(), \
rot_dim, q_pe_stride_token, q_pe_stride_head, \
k_pe_stride, kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>( \
kv_cache.mutable_data_ptr()), \
slot_mapping.const_data_ptr<int64_t>(), \
block_stride, entry_stride, kv_lora_rank, \
block_size, \
kv_cache_quant_scale.const_data_ptr<float>()); \
} \
}); \
}); \
} while (false)
// Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -208,64 +215,69 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
// Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and
// concat_and_cache_mla.
void concat_and_cache_mla_rope_fused(
torch::Tensor& positions, // [num_tokens]
torch::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim]
torch::Tensor& k_pe, // [num_tokens, rot_dim]
torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
torch::stable::Tensor& positions, // [num_tokens]
torch::stable::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim]
torch::stable::Tensor& k_pe, // [num_tokens, rot_dim]
torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank]
torch::stable::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
bool rope_is_neox,
torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
torch::Tensor&
torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
torch::stable::Tensor&
kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
const std::string& kv_cache_dtype,
torch::stable::Tensor& kv_cache_quant_scale) {
// NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
// slot_mapping.size(0) because of padding for CUDA graphs.
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
// both include padding.
// In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
// since key includes padding for CUDA graphs, while slot_mapping does not.
// In this case, slot_mapping.size(0) represents the actual number of tokens
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0)
// because both include padding.
// In vLLM V1, however, key.size(0) can be larger than
// slot_mapping.size(0) since key includes padding for CUDA graphs,
// while slot_mapping does not. In this case,
// slot_mapping.size(0) represents the actual number of tokens
// before padding.
// For compatibility with both cases, we use slot_mapping.size(0) as the
// number of tokens.
int num_tokens = slot_mapping.size(0);
int num_padded_tokens = q_pe.size(0);
TORCH_CHECK_GE(num_padded_tokens, num_tokens);
// For compatibility with both cases, we use slot_mapping.size(0) as
// the number of tokens.
const int64_t num_tokens = slot_mapping.size(0);
const int64_t num_padded_tokens = q_pe.size(0);
STD_TORCH_CHECK(num_padded_tokens >= num_tokens);
const int num_q_heads = q_pe.size(1);
const int rot_dim = q_pe.size(2);
const int kv_lora_rank = kv_c.size(1);
TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
TORCH_CHECK_EQ(positions.dim(), 1);
TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
STD_TORCH_CHECK(positions.size(0) == num_padded_tokens);
STD_TORCH_CHECK(positions.dim() == 1);
STD_TORCH_CHECK(positions.scalar_type() ==
torch::headeronly::ScalarType::Long);
TORCH_CHECK_EQ(q_pe.dim(), 3);
TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
STD_TORCH_CHECK(q_pe.dim() == 3);
STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens);
STD_TORCH_CHECK(q_pe.size(1) == num_q_heads);
STD_TORCH_CHECK(q_pe.size(2) == rot_dim);
TORCH_CHECK_EQ(k_pe.dim(), 2);
TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
STD_TORCH_CHECK(k_pe.dim() == 2);
STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens);
STD_TORCH_CHECK(k_pe.size(1) == rot_dim);
STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type());
TORCH_CHECK_EQ(kv_c.dim(), 2);
TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
STD_TORCH_CHECK(kv_c.dim() == 2);
STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens);
STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank);
STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type());
TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim);
STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type());
TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens);
STD_TORCH_CHECK(slot_mapping.scalar_type() ==
torch::headeronly::ScalarType::Long);
TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
TORCH_CHECK_EQ(kv_cache.dim(), 3);
STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim);
STD_TORCH_CHECK(kv_cache.dim() == 3);
TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1);
TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float);
STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1);
STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() ==
torch::headeronly::ScalarType::Float);
int64_t q_pe_stride_token = q_pe.stride(0);
int64_t q_pe_stride_head = q_pe.stride(1);
@@ -286,9 +298,10 @@ void concat_and_cache_mla_rope_fused(
dim3 grid(num_tokens, 1, 1);
dim3 block(thread_block_size, 1, 1);
const at::cuda::OptionalCUDAGuard device_guard(device_of(positions));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const torch::stable::accelerator::DeviceGuard device_guard(
positions.get_device_index());
const cudaStream_t stream = get_current_cuda_stream();
DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED);
}
csrc/concat_mla_q.cuh → csrc/libtorch_stable/concat_mla_q.cuh
View File
csrc/cuda_vec_utils.cuh → csrc/libtorch_stable/cuda_vec_utils.cuh
View File
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp → csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp
View File
csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
+1 -1
View File
@@ -2,7 +2,7 @@
#include <torch/csrc/stable/tensor.h>
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
#include "broadcast_load_epilogue_c2x.hpp"
/*
This file defines custom epilogues for fusing channel scales, token scales,
csrc/libtorch_stable/fp32_router_gemm.cu
+223
View File
@@ -0,0 +1,223 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
// Supports bf16 or fp32 activation; weight is always fp32.
// Adapted from dsv3_router_gemm_float_out.cu.
#include <cuda_bf16.h>
#include <cuda_runtime.h>
// ---------------------------------------------------------------------------
// Load helpers
// ---------------------------------------------------------------------------
// Load VPT fp32 values from the weight matrix (always fp32).
// VPT=4 when activation is fp32 (one float4 load)
// VPT=8 when activation is bf16 (two float4 loads)
template <int VPT>
__device__ __forceinline__ void load_weight(float const* ptr, float* dst);
template <>
__device__ __forceinline__ void load_weight<4>(float const* ptr, float* dst) {
float4 v = *reinterpret_cast<float4 const*>(ptr);
dst[0] = v.x;
dst[1] = v.y;
dst[2] = v.z;
dst[3] = v.w;
}
template <>
__device__ __forceinline__ void load_weight<8>(float const* ptr, float* dst) {
float4 v0 = *reinterpret_cast<float4 const*>(ptr);
float4 v1 = *reinterpret_cast<float4 const*>(ptr + 4);
dst[0] = v0.x;
dst[1] = v0.y;
dst[2] = v0.z;
dst[3] = v0.w;
dst[4] = v1.x;
dst[5] = v1.y;
dst[6] = v1.z;
dst[7] = v1.w;
}
// Load VPT activation values and convert to fp32.
template <typename T, int VPT>
__device__ __forceinline__ void load_activation(T const* ptr, float* dst);
// fp32 activation: one float4 load, no conversion needed.
template <>
__device__ __forceinline__ void load_activation<float, 4>(float const* ptr,
float* dst) {
float4 v = *reinterpret_cast<float4 const*>(ptr);
dst[0] = v.x;
dst[1] = v.y;
dst[2] = v.z;
dst[3] = v.w;
}
// bf16 activation: one uint4 load (8 × bf16) + element-wise conversion.
template <>
__device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
__nv_bfloat16 const* ptr, float* dst) {
uint4 v = *reinterpret_cast<uint4 const*>(ptr);
__nv_bfloat16 const* bf16_ptr = reinterpret_cast<__nv_bfloat16 const*>(&v);
#pragma unroll
for (int i = 0; i < 8; i++) dst[i] = __bfloat162float(bf16_ptr[i]);
}
// ---------------------------------------------------------------------------
// Kernel
// ---------------------------------------------------------------------------
// InputT : type of activation (float or __nv_bfloat16)
// Weight is always fp32; output is always fp32.
// VPT = 16 / sizeof(InputT): 4 for fp32, 8 for bf16
template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
int kHiddenDim>
__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
float* out, InputT const* mat_a, float const* mat_b) {
constexpr int VPT = 16 / sizeof(InputT);
constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
constexpr int kWarpSize = 32;
constexpr int kNumWarps = kBlockSize / kWarpSize;
int const n_idx = blockIdx.x;
int const tid = threadIdx.x;
int const warpId = tid / kWarpSize;
int const laneId = tid % kWarpSize;
float acc[kNumTokens] = {};
__shared__ float sm_reduction[kNumTokens][kNumWarps];
float const* b_col = mat_b + n_idx * kHiddenDim;
int k_bases[k_iterations];
#pragma unroll
for (int ki = 0; ki < k_iterations; ki++) {
k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
}
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
asm volatile("griddepcontrol.wait;");
#endif
for (int ki = 0; ki < k_iterations; ki++) {
int const k_base = k_bases[ki];
float b_float[VPT];
load_weight<VPT>(b_col + k_base, b_float);
#pragma unroll
for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
float a_float[VPT];
load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
a_float);
#pragma unroll
for (int k = 0; k < VPT; k++) {
acc[m_idx] += a_float[k] * b_float[k];
}
}
}
// Warp-level butterfly reduction
#pragma unroll
for (int m = 0; m < kNumTokens; m++) {
float sum = acc[m];
sum += __shfl_xor_sync(0xffffffff, sum, 16);
sum += __shfl_xor_sync(0xffffffff, sum, 8);
sum += __shfl_xor_sync(0xffffffff, sum, 4);
sum += __shfl_xor_sync(0xffffffff, sum, 2);
sum += __shfl_xor_sync(0xffffffff, sum, 1);
if (laneId == 0) sm_reduction[m][warpId] = sum;
}
__syncthreads();
if (tid == 0) {
#pragma unroll
for (int m = 0; m < kNumTokens; m++) {
float final_sum = 0.0f;
#pragma unroll
for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
out[m * kNumExperts + n_idx] = final_sum;
}
}
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
asm volatile("griddepcontrol.launch_dependents;");
#endif
}
// ---------------------------------------------------------------------------
// Launcher
// ---------------------------------------------------------------------------
template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
void invokeFp32RouterGemm(float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream) {
constexpr int kBlockSize = 128;
cudaLaunchConfig_t config;
config.gridDim = kNumExperts;
config.blockDim = kBlockSize;
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.numAttrs = 1;
config.attrs = attrs;
cudaLaunchKernelEx(&config,
fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
kNumExperts, kHiddenDim>,
output, mat_a, mat_b);
}
// ---------------------------------------------------------------------------
// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
// ---------------------------------------------------------------------------
#define INSTANTIATE(T, M) \
template void invokeFp32RouterGemm<T, M, 256, 3072>( \
float*, T const*, float const*, cudaStream_t);
#define INSTANTIATE_ALL(T) \
INSTANTIATE(T, 1) \
INSTANTIATE(T, 2) \
INSTANTIATE(T, 3) \
INSTANTIATE(T, 4) \
INSTANTIATE(T, 5) \
INSTANTIATE(T, 6) \
INSTANTIATE(T, 7) \
INSTANTIATE(T, 8) \
INSTANTIATE(T, 9) \
INSTANTIATE(T, 10) \
INSTANTIATE(T, 11) \
INSTANTIATE(T, 12) \
INSTANTIATE(T, 13) \
INSTANTIATE(T, 14) \
INSTANTIATE(T, 15) \
INSTANTIATE(T, 16) \
INSTANTIATE(T, 17) \
INSTANTIATE(T, 18) \
INSTANTIATE(T, 19) \
INSTANTIATE(T, 20) \
INSTANTIATE(T, 21) \
INSTANTIATE(T, 22) \
INSTANTIATE(T, 23) \
INSTANTIATE(T, 24) \
INSTANTIATE(T, 25) \
INSTANTIATE(T, 26) \
INSTANTIATE(T, 27) \
INSTANTIATE(T, 28) \
INSTANTIATE(T, 29) \
INSTANTIATE(T, 30) \
INSTANTIATE(T, 31) \
INSTANTIATE(T, 32)
INSTANTIATE_ALL(float)
INSTANTIATE_ALL(__nv_bfloat16)
#undef INSTANTIATE_ALL
#undef INSTANTIATE
csrc/libtorch_stable/fp32_router_gemm_entry.cu
+127
View File
@@ -0,0 +1,127 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include "core/registration.h"
#include "libtorch_stable/torch_utils.h"
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <stdexcept>
namespace {
inline int getSMVersion() {
auto* props = get_device_prop();
return props->major * 10 + props->minor;
}
} // namespace
static constexpr int FP32_NUM_EXPERTS = 256;
static constexpr int FP32_HIDDEN_DIM = 3072;
static constexpr int FP32_MAX_TOKENS = 32;
// Forward declarations — 4 template params must match fp32_router_gemm.cu
template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
void invokeFp32RouterGemm(float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream);
// LoopUnroller templated on InputT
template <typename InputT, int kBegin, int kEnd>
struct Fp32LoopUnroller {
static void unroll(int num_tokens, float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream) {
if (num_tokens == kBegin) {
invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
output, mat_a, mat_b, stream);
} else {
Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
mat_a, mat_b, stream);
}
}
};
template <typename InputT, int kEnd>
struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
static void unroll(int num_tokens, float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream) {
if (num_tokens == kEnd) {
invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
output, mat_a, mat_b, stream);
} else {
throw std::invalid_argument(
"fp32_router_gemm: num_tokens must be in [1, 32]");
}
}
};
void fp32_router_gemm(
torch::stable::Tensor& output, // [num_tokens, num_experts]
torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
torch::stable::Tensor const& mat_b // [num_experts, hidden_dim]
) {
STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
"fp32_router_gemm: all tensors must be CUDA tensors");
STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
output.get_device_index() == mat_b.get_device_index(),
"fp32_router_gemm: all tensors must be on the same device");
STD_TORCH_CHECK(
output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(),
"fp32_router_gemm: all tensors must be contiguous");
const int num_tokens = mat_a.size(0);
const int num_experts = mat_b.size(0);
const int hidden_dim = mat_a.size(1);
STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts,
"fp32_router_gemm: output must have shape [num_tokens, "
"num_experts]");
STD_TORCH_CHECK(
mat_a.size(1) == mat_b.size(1),
"fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
"fp32_router_gemm: expected hidden_dim=3072");
STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
"fp32_router_gemm: expected num_experts=256");
STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
"fp32_router_gemm: num_tokens must be in [0, 32]");
STD_TORCH_CHECK(
mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"fp32_router_gemm: mat_a must be float32 or bfloat16");
STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float,
"fp32_router_gemm: mat_b (weight) must be float32");
STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float,
"fp32_router_gemm: output must be float32");
if (num_tokens == 0) {
return;
}
STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
auto stream = get_current_cuda_stream(mat_a.get_device_index());
float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
float const* mat_b_ptr = reinterpret_cast<float const*>(mat_b.data_ptr());
if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
auto const* mat_a_ptr =
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
} else {
auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
}
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm));
}
csrc/libtorch_stable/fused_qknorm_rope_kernel.cu
+1 -1
View File
@@ -20,7 +20,7 @@
#include "torch_utils.h"
#include "../async_util.cuh"
#include "async_util.cuh"
#include "../cuda_compat.h"
#include "../type_convert.cuh"
#include "dispatch_utils.h"
csrc/launch_bounds_utils.h → csrc/libtorch_stable/launch_bounds_utils.h
View File
csrc/libtorch_stable/layernorm_kernels.cu
+4 -6
View File
@@ -78,8 +78,7 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
}
v_out[i] = dst;
}
@@ -143,8 +142,7 @@ fused_add_rms_norm_kernel(
#pragma unroll
for (int j = 0; j < width; ++j) {
float x = Converter::convert(res.data[j]);
float wf = Converter::convert(w.data[j]);
out.data[j] = Converter::convert(x * s_variance * wf);
out.data[j] = Converter::convert(x * s_variance) * w.data[j];
}
input_v[strided_id] = out;
}
@@ -183,8 +181,8 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
input[blockIdx.x * input_stride + idx] =
(scalar_t)(x * s_variance) * weight[idx];
}
}
csrc/libtorch_stable/layernorm_quant_kernels.cu
+6 -17
View File
@@ -66,13 +66,8 @@ __global__ void rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
float w = static_cast<float>(src2.val[j]);
// Round normalized result through scalar_t to match the precision of the
// unfused composite (rms_norm writes scalar_t, then
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
// Without this round, the fused path is strictly more accurate and
// disagrees with the composite at exact E4M3 quantization tie boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
// Multiply in weight's native dtype to match rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
@@ -142,12 +137,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(res.data[i]);
float wf = Converter::convert(w.data[i]);
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries. We use the
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
HipT out_norm_h = Converter::convert(x * s_variance * wf);
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
}
@@ -192,10 +183,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
float w = (float)weight[idx];
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
static_cast<float>(out_norm), scale_inv);
}
csrc/nvfp4_kv_cache_kernels.cu → csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
+36 -34
View File
@@ -17,11 +17,8 @@
#define NVFP4_ENABLE_ELTS16 1
#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "dispatch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
namespace vllm {
@@ -184,12 +181,13 @@ __global__ void reshape_and_cache_nvfp4_kernel(
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
// Each KV side contains both data and scale:
// page = [K_data | K_scale | V_data | V_scale]
void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
torch::Tensor& k_scale,
torch::Tensor& v_scale) {
void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
torch::stable::Tensor& value,
torch::stable::Tensor& key_cache,
torch::stable::Tensor& value_cache,
torch::stable::Tensor& slot_mapping,
torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale) {
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
@@ -200,17 +198,18 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
// key_cache is kv_cache[:, 0] with shape
// [num_blocks, block_size, num_heads, full_dim] in logical order.
// Strides encode the physical layout (HND or NHD).
TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
TORCH_CHECK(key_cache.size(3) == full_dim,
"key_cache last dim must be data_dim + scale_dim, got ",
key_cache.size(3), " expected ", full_dim);
STD_TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
STD_TORCH_CHECK(key_cache.size(3) == full_dim,
"key_cache last dim must be data_dim + scale_dim, got ",
key_cache.size(3), " expected ", full_dim);
int block_size = key_cache.size(1);
TORCH_CHECK(head_size % 16 == 0,
"head_size must be divisible by 16 for NVFP4 KV cache");
TORCH_CHECK(block_size % 4 == 0,
"block_size must be divisible by 4 for NVFP4 KV cache swizzle");
STD_TORCH_CHECK(head_size % 16 == 0,
"head_size must be divisible by 16 for NVFP4 KV cache");
STD_TORCH_CHECK(block_size % 4 == 0,
"block_size must be divisible by 4 for NVFP4 KV cache "
"swizzle");
// Detect physical layout from strides (based on full_dim).
// HND: head stride > block_offset stride.
@@ -230,8 +229,9 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
// Scale follows data within each KV side.
int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
uint8_t* key_scale_ptr = key_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
uint8_t* value_scale_ptr =
value_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
// Scale strides: same page stride, inner strides from layout.
int64_t scale_block_stride = data_block_stride;
@@ -244,8 +244,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
scale_block_offset_stride = (int64_t)num_heads * scale_dim;
}
const float* k_scale_ptr = k_scale.data_ptr<float>();
const float* v_scale_ptr = v_scale.data_ptr<float>();
const float* k_scale_ptr = k_scale.const_data_ptr<float>();
const float* v_scale_ptr = v_scale.const_data_ptr<float>();
int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
int total_groups = num_heads * groups_per_head;
@@ -256,20 +256,22 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
dim3 grid(num_tokens);
dim3 block(num_threads);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const torch::stable::accelerator::DeviceGuard device_guard(
key.get_device_index());
const cudaStream_t stream = get_current_cuda_stream();
AT_DISPATCH_REDUCED_FLOATING_TYPES(
VLLM_STABLE_DISPATCH_HALF_TYPES(
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
<<<grid, block, 0, stream>>>(
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
key_scale_ptr, value_scale_ptr,
slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
key.stride(0), value.stride(0), num_heads, head_size,
block_size, data_block_stride, data_head_stride,
data_block_offset_stride, scale_block_stride, scale_head_stride,
scale_block_offset_stride);
key.const_data_ptr<scalar_t>(),
value.const_data_ptr<scalar_t>(),
key_cache.mutable_data_ptr<uint8_t>(),
value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
num_heads, head_size, block_size, data_block_stride,
data_head_stride, data_block_offset_stride, scale_block_stride,
scale_head_stride, scale_block_offset_stride);
});
}
csrc/libtorch_stable/ops.h
+133 -4
View File
@@ -3,10 +3,6 @@
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#ifndef USE_ROCM
torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
torch::stable::Tensor const& perm);
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
torch::stable::Tensor& output_s,
@@ -28,6 +24,10 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
int64_t group_size, double eps, double int8_min,
double int8_max);
#ifndef USE_ROCM
torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
torch::stable::Tensor const& perm);
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -355,3 +355,132 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
int64_t tokens);
int64_t ggml_moe_get_block_size(int64_t type);
void paged_attention_v1(
torch::stable::Tensor& out, torch::stable::Tensor& query,
torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
const std::optional<torch::stable::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale, const int64_t tp_rank,
const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step);
void paged_attention_v2(
torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
int64_t block_size, int64_t max_seq_len,
const std::optional<torch::stable::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale, const int64_t tp_rank,
const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step);
// Cache ops (shared CUDA/ROCm)
void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
int64_t block_size_in_bytes,
const torch::stable::Tensor& block_mapping);
// Batch swap: submit all block copies in a single driver call.
void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
const torch::stable::Tensor& dst_ptrs,
const torch::stable::Tensor& sizes,
bool is_src_access_order_any);
void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
torch::stable::Tensor& key_cache,
torch::stable::Tensor& value_cache,
torch::stable::Tensor& slot_mapping,
const std::string& kv_cache_dtype,
torch::stable::Tensor& k_scale,
torch::stable::Tensor& v_scale);
void reshape_and_cache_flash(
torch::stable::Tensor& key, torch::stable::Tensor& value,
torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype,
torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale);
void concat_and_cache_mla(torch::stable::Tensor& kv_c,
torch::stable::Tensor& k_pe,
torch::stable::Tensor& kv_cache,
torch::stable::Tensor& slot_mapping,
const std::string& kv_cache_dtype,
torch::stable::Tensor& scale);
// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
void concat_and_cache_mla_rope_fused(
torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c,
torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox,
torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache,
const std::string& kv_cache_dtype,
torch::stable::Tensor& kv_cache_quant_scale);
// Just for unittest
void convert_fp8(torch::stable::Tensor& dst_cache,
torch::stable::Tensor& src_cache, const double scale,
const std::string& kv_cache_dtype);
void gather_and_maybe_dequant_cache(
torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
// ENTRIES...]
torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
int64_t num_tokens, const std::string& kv_cache_dtype,
torch::stable::Tensor const& scale,
std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
// TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
void cp_gather_cache(
torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
// ENTRIES...]
torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
int64_t batch_size,
std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
// Gather and upconvert FP8 KV cache to BF16 workspace
void cp_gather_and_upconvert_fp8_kv_cache(
torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
// 656]
torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
torch::stable::Tensor const& seq_lens, // [BATCH]
torch::stable::Tensor const& workspace_starts, // [BATCH]
int64_t batch_size);
// Indexer K quantization and cache function
void indexer_k_quant_and_cache(
torch::stable::Tensor& k, // [num_tokens, head_dim]
torch::stable::Tensor& kv_cache, // [num_blocks, block_size,
// cache_stride]
torch::stable::Tensor& slot_mapping, // [num_tokens]
int64_t quant_block_size, // quantization block size
const std::string& scale_fmt);
// Concatenate query nope and rope for MLA/DSA attention
void concat_mla_q(
torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim +
// rope_dim]
// Extract function to gather quantized K cache
void cp_gather_indexer_k_quant_cache(
const torch::stable::Tensor& kv_cache, // [num_blocks, block_size,
// cache_stride]
torch::stable::Tensor& dst_k, // [num_tokens, head_dim]
torch::stable::Tensor& dst_scale, // [num_tokens, head_dim /
// quant_block_size * 4]
const torch::stable::Tensor& block_table, // [batch_size, num_blocks]
const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1]
csrc/persistent_topk.cuh → csrc/libtorch_stable/persistent_topk.cuh
View File
csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+2 -2
View File
@@ -17,7 +17,7 @@
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "../../cuda_vec_utils.cuh"
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
@@ -25,7 +25,7 @@
#include <cuda_fp8.h>
#include "cuda_utils.h"
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
// Define before including nvfp4_utils.cuh so the header
// can use this macro during compilation.
csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu
+2 -2
View File
@@ -27,14 +27,14 @@
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "../../cuda_vec_utils.cuh"
#include "cuda_utils.h"
#include "nvfp4_utils.cuh"
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
namespace vllm {
csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu
+2 -2
View File
@@ -17,7 +17,7 @@
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "../../cuda_vec_utils.cuh"
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
@@ -26,7 +26,7 @@
#include "cuda_utils.h"
#include "nvfp4_utils.cuh"
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
namespace vllm {
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+2 -2
View File
@@ -23,10 +23,10 @@
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "../../cuda_vec_utils.cuh"
#include "cuda_utils.h"
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
// Define before including nvfp4_utils.cuh so the header
// can use this macro during compilation.
csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh
+1 -1
View File
@@ -20,7 +20,7 @@
#include <cuda_fp8.h>
#include <utility>
#include "cuda_vec_utils.cuh"
#include "../../cuda_vec_utils.cuh"
#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
CUDA_VERSION >= 12090
csrc/quantization/gguf/dequantize.cuh → csrc/libtorch_stable/quantization/gguf/dequantize.cuh
View File
csrc/quantization/gguf/ggml-common.h → csrc/libtorch_stable/quantization/gguf/ggml-common.h
View File
csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu
+5 -8
View File
@@ -7,14 +7,11 @@
#include <torch/csrc/stable/ops.h>
// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
// moved to libtorch_stable) to avoid unnecessary reformatting that would break
// git rename detection and pollute blame history.
#include "../../../quantization/gguf/ggml-common.h"
#include "../../../quantization/gguf/vecdotq.cuh"
#include "../../../quantization/gguf/dequantize.cuh"
#include "../../../quantization/gguf/mmvq.cuh"
#include "../../../quantization/gguf/mmq.cuh"
#include "ggml-common.h"
#include "vecdotq.cuh"
#include "dequantize.cuh"
#include "mmvq.cuh"
#include "mmq.cuh"
#include "moe.cuh"
#include "moe_vec.cuh"
csrc/quantization/gguf/mmq.cuh → csrc/libtorch_stable/quantization/gguf/mmq.cuh
View File
csrc/quantization/gguf/mmvq.cuh → csrc/libtorch_stable/quantization/gguf/mmvq.cuh
View File
csrc/quantization/gguf/vecdotq.cuh → csrc/libtorch_stable/quantization/gguf/vecdotq.cuh
View File
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+33 -11
View File
@@ -7,7 +7,11 @@
#include <cmath>
#include <cuda_fp8.h>
#ifdef USE_ROCM
#include <hip/hip_fp8.h>
#else
#include <cuda_fp8.h>
#endif
#include "libtorch_stable/quantization/vectorization.cuh"
#include "libtorch_stable/quantization/vectorization_utils.cuh"
@@ -15,12 +19,23 @@
#include "libtorch_stable/torch_utils.h"
__device__ __forceinline__ float GroupReduceMax(float val) {
#ifdef USE_ROCM
// 16-thread logical groups may pack up to four per 64-lane wavefront; use a
// 64-bit mask and explicit width so shuffles stay within each group.
const int lane_in_wave = threadIdx.x % warpSize;
const unsigned long long mask = 0xFFFFull << ((lane_in_wave / 16) * 16);
val = fmaxf(val, __shfl_xor_sync(mask, val, 8, 16));
val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 16));
val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 16));
val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 16));
#else
unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
#endif
return val;
}
@@ -237,10 +252,12 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "per_token_group_quant_8bit", ([&] {
if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);
} else if (dst_type == torch::headeronly::ScalarType::Char) {
if (dst_type == torch::headeronly::ScalarType::Char) {
LAUNCH_KERNEL(scalar_t, int8_t);
} else {
VLLM_STABLE_DISPATCH_FP8_TYPES(
dst_type, "per_token_group_quant_8bit_fp8",
([&] { LAUNCH_KERNEL(scalar_t, fp8_t); }));
}
}));
@@ -317,10 +334,18 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
// 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
// 8 lanes within the warp that share a group.
#ifdef USE_ROCM
const int lane_in_wave = threadIdx.x % warpSize;
const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));
#else
unsigned mask = 0xffu << (threadIdx.x & 24u);
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
#endif
float y_s = local_absmax / max_8bit;
y_s = fmaxf(y_s, 1e-10f);
@@ -503,15 +528,12 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
} else if (dst_type == torch::headeronly::ScalarType::Char) {
if (dst_type == torch::headeronly::ScalarType::Char) {
LAUNCH_REG_KERNEL(scalar_t, int8_t);
} else {
STD_TORCH_CHECK(
false,
"per_token_group_quant_8bit_packed only supports FP8/INT8 "
"outputs.");
VLLM_STABLE_DISPATCH_FP8_TYPES(
dst_type, "per_token_group_quant_8bit_packed_fp8",
([&] { LAUNCH_REG_KERNEL(scalar_t, fp8_t); }));
}
}));
csrc/libtorch_stable/topk.cu
+1 -1
View File
@@ -7,7 +7,7 @@
#include "torch_utils.h"
#ifndef USE_ROCM
#include "../persistent_topk.cuh"
#include "persistent_topk.cuh"
#endif
namespace {
csrc/libtorch_stable/torch_bindings.cpp
+155 -10
View File
@@ -7,11 +7,6 @@
// Note: We register under namespace "_C" so ops are accessible as
// torch.ops._C.<op_name> for compatibility with existing code.
STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
#ifndef USE_ROCM
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
#endif
#ifndef USE_ROCM
// Compute per-token-group FP8 quantized tensor and scaling factor.
// The dummy arguments are here so we can correctly fuse with RMSNorm.
ops.def(
@@ -32,6 +27,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
"()");
#ifndef USE_ROCM
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
#endif
#ifndef USE_ROCM
// CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
// quantization, as well as bias
ops.def(
@@ -247,6 +247,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def(
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
// BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
// conditionally compiled so impl registration is in source file
ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
ops.def(
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
@@ -474,14 +478,36 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor? initial_state_idx,"
"Tensor? cu_chunk_seqlen,"
"Tensor? last_chunk_indices) -> ()");
// Attention ops
// Compute the attention between an input query and the cached
// keys/values using PagedAttention.
ops.def(
"paged_attention_v1("
" Tensor! out, Tensor query, Tensor key_cache,"
" Tensor value_cache, int num_kv_heads, float scale,"
" Tensor block_tables, Tensor seq_lens, int block_size,"
" int max_seq_len, Tensor? alibi_slopes,"
" str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
" int tp_rank, int blocksparse_local_blocks,"
" int blocksparse_vert_stride, int blocksparse_block_size,"
" int blocksparse_head_sliding_step) -> ()");
// PagedAttention V2.
ops.def(
"paged_attention_v2("
" Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
" Tensor! tmp_out, Tensor query, Tensor key_cache,"
" Tensor value_cache, int num_kv_heads, float scale,"
" Tensor block_tables, Tensor seq_lens, int block_size,"
" int max_seq_len, Tensor? alibi_slopes,"
" str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
" int tp_rank, int blocksparse_local_blocks,"
" int blocksparse_vert_stride, int blocksparse_block_size,"
" int blocksparse_head_sliding_step) -> ()");
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
#ifndef USE_ROCM
ops.impl("permute_cols", TORCH_BOX(&permute_cols));
#endif
#ifndef USE_ROCM
// Per-token group quantization
ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));
ops.impl("per_token_group_fp8_quant_packed",
@@ -489,6 +515,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("per_token_group_quant_int8",
TORCH_BOX(&per_token_group_quant_int8));
#ifndef USE_ROCM
ops.impl("permute_cols", TORCH_BOX(&permute_cols));
#endif
#ifndef USE_ROCM
// CUTLASS scaled_mm ops
ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm));
ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp));
@@ -581,6 +612,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
}
// These capability-check functions take only primitive args (no tensors), so
@@ -603,4 +637,115 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
}
// Cache ops
STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
// Swap in (out) the cache blocks from src to dst.
ops.def(
"swap_blocks(Tensor src, Tensor! dst,"
" int block_size_in_bytes, Tensor block_mapping) -> ()");
// Batch swap: submit all block copies in a single driver call.
ops.def(
"swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
" Tensor sizes,"
" bool is_src_access_order_any=False) -> ()");
// Reshape the key and value tensors and cache them.
ops.def(
"reshape_and_cache(Tensor key, Tensor value,"
" Tensor! key_cache, Tensor! value_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale) -> ()");
// Reshape the key and value tensors and cache them.
ops.def(
"reshape_and_cache_flash(Tensor key, Tensor value,"
" Tensor! key_cache,"
" Tensor! value_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale) -> ()");
// Concat kv_c and k_pe and cache them.
ops.def(
"concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
" Tensor! kv_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor scale) -> ()");
// Rotate Q and K, then write to kv cache for MLA
ops.def(
"concat_and_cache_mla_rope_fused("
" Tensor positions,"
" Tensor! q_pe,"
" Tensor! k_pe,"
" Tensor kv_c,"
" Tensor cos_sin_cache,"
" bool is_neox,"
" Tensor slot_mapping,"
" Tensor! kv_cache,"
" str kv_cache_dtype,"
" Tensor kv_cache_scale) -> ()");
// Convert the key and value cache to fp8 data type.
ops.def(
"convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
"str kv_cache_dtype) -> ()");
// Gather cache blocks from src_cache to dst, dequantizing from
// src_cache's dtype to dst's dtype if necessary.
ops.def(
"gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
" Tensor block_table, Tensor cu_seq_lens, "
" Tensor token_to_seq, "
" int num_tokens, "
" str kv_cache_dtype, "
" Tensor scale, Tensor? seq_starts) -> ()");
ops.def(
"cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
"Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
ops.def(
"cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
"Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
"batch_size) -> ()");
ops.def(
"indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
"slot_mapping, "
"int quant_block_size, str kv_cache_dtype) -> ()");
ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
ops.def(
"cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
"dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
}
STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
}
STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
ops.impl("swap_blocks", TORCH_BOX(&swap_blocks));
ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
ops.impl("concat_and_cache_mla_rope_fused",
TORCH_BOX(&concat_and_cache_mla_rope_fused));
ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
ops.impl("gather_and_maybe_dequant_cache",
TORCH_BOX(&gather_and_maybe_dequant_cache));
ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache));
ops.impl("cp_gather_and_upconvert_fp8_kv_cache",
TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache));
ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache));
ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q));
ops.impl("cp_gather_indexer_k_quant_cache",
TORCH_BOX(&cp_gather_indexer_k_quant_cache));
}
REGISTER_EXTENSION(_C_stable_libtorch)
csrc/libtorch_stable/torch_utils.h
+1 -5
View File
@@ -6,11 +6,7 @@
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/shim_utils.h>
#ifndef USE_ROCM
#include <cuda_runtime.h>
#else
#include <hip/hip_runtime.h>
#endif
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <deque>
csrc/ops.h
-23
View File
@@ -31,29 +31,6 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
return new_tensor;
}
void paged_attention_v1(
torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale, const int64_t tp_rank,
const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step);
void paged_attention_v2(
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale, const int64_t tp_rank,
const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step);
// rms_norm and fused_add_rms_norm declarations also exist in
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
// because the CPU build still uses these torch::Tensor declarations.
csrc/quantization/w8a8/fp8/amd/quant_utils.cuh
+12 -9
View File
@@ -6,6 +6,7 @@
#include <hip/hip_bfloat16.h>
#include "../../../../attention/attention_dtypes.h"
#include <torch/headeronly/core/ScalarType.h>
namespace vllm {
#ifdef USE_ROCM
@@ -642,27 +643,29 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \
vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \
if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
if (SRC_DTYPE == at::ScalarType::Float) { \
if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \
} else if (SRC_DTYPE == at::ScalarType::Half) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \
} else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \
} else { \
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
STD_TORCH_CHECK(false, \
"Unsupported input type of kv cache: ", SRC_DTYPE); \
} \
} else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \
if (SRC_DTYPE == at::ScalarType::Float) { \
if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
} else if (SRC_DTYPE == at::ScalarType::Half) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
} else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
} else { \
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
STD_TORCH_CHECK(false, \
"Unsupported input type of kv cache: ", SRC_DTYPE); \
} \
} else { \
TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
}
} // namespace fp8
csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
+17 -13
View File
@@ -1,6 +1,7 @@
#pragma once
#include "../../../../attention/attention_dtypes.h"
#include <torch/headeronly/core/ScalarType.h>
#include <assert.h>
#include <float.h>
#include <stdint.h>
@@ -546,37 +547,40 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \
vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \
if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
if (SRC_DTYPE == at::ScalarType::Float) { \
if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \
} else if (SRC_DTYPE == at::ScalarType::Half) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \
} else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \
} else { \
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
STD_TORCH_CHECK(false, \
"Unsupported input type of kv cache: ", SRC_DTYPE); \
} \
} else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \
if (SRC_DTYPE == at::ScalarType::Float) { \
if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
} else if (SRC_DTYPE == at::ScalarType::Half) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
} else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
} else { \
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
STD_TORCH_CHECK(false, \
"Unsupported input type of kv cache: ", SRC_DTYPE); \
} \
} else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \
if (SRC_DTYPE == at::ScalarType::Float) { \
if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
} else if (SRC_DTYPE == at::ScalarType::Half) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
} else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
} else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
} else { \
TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
STD_TORCH_CHECK(false, \
"Unsupported input type of kv cache: ", SRC_DTYPE); \
} \
} else { \
TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
}
} // namespace fp8
csrc/rocm/ops.h
+9
View File
@@ -18,6 +18,15 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
const at::Tensor& scale_a, const at::Tensor& scale_b,
const int64_t CuCount);
torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros, torch::Tensor b_scales,
torch::Tensor b_g_idx, bool use_v2_format);
torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros,
torch::Tensor b_scales,
torch::Tensor b_g_idx, bool use_v2_format);
void paged_attention(
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
csrc/rocm/q_gemm_rdna3.cu
+780
View File
@@ -0,0 +1,780 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the
// activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit
// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes:
//
// 1. Direct write to the T-typed output via packed CAS-loop on a 64-bit
// word (atomic_add_pk4_{f16,bf16}). gfx11 has no native
// v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with
// global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch
// buffer + memset + cast-pass that an fp32-accumulator design would
// need; the caller passes a zero-initialised T-typed output tensor
// and every block atomically adds its partial sum into it.
//
// 2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only
// "upper nibble * 16" trick, which would overflow the 7-bit bf16
// mantissa. See qdq_4_rdna3.cuh for details.
//
// 3. Wave32 geometry sized for high CU saturation: THREADS_X=256
// (8 waves per block) and BLOCK_KN_SIZE=256, with each thread
// computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE
// splits K and the output is atomically accumulated. fp16 uses
// v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot;
// bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates
// with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch
// based on size_m.
//
// 4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in
// q_gemm_rdna3_wmma.cu (separate translation unit) where
// v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays
// scalar (the bit-trick dequant beats WMMA below M=64).
#include <cstdint>
#include <cstdio>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include "qdq_4_rdna3.cuh"
#if defined(__HIPCC__) && defined(__gfx1100__)
#define __HIP__RDNA3__
#endif
namespace vllm {
namespace gptq_rdna3 {
// BLOCK_KN_SIZE = 256 (was 128 in exllama). Each block covers 256 K
// elements and THREADS_X*4 = 1024 N columns. For Qwen-class K=4096 this
// halves gridDim.z (32 → 16) and therefore halves the atomic count per
// output position vs the exllama default. THREADS_X=256 = 8 waves on RDNA3
// wave32; with ~32 wave slots per CU we still fit 4 blocks per CU at peak.
//
// We tried BLOCK_KN_SIZE=512 (microbench on Qwen3.6-27B): bf16 improved
// 5-10% at large M (atomic CAS halved), but fp16 decode regressed up to
// +40% on qkv-square (32 → 45 μs at M=1). Cause: 16 waves/block × 16
// total blocks for [M=1, K=N=4096] only saturates ~8 of the 96 CUs,
// breaking memory-latency hiding for the fp16 path which is already
// memory-bound. Reverted to 256; bf16 keeps most of its gains from the
// fp32 dequant rewrite alone.
#define BLOCK_KN_SIZE 256
#define THREADS_X 256
// Device code below is RDNA3-only; non-RDNA3 device passes fall through to
// the empty __global__ stub at the #else below for symbol parity.
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
// ---------------------------------------------------------------------------
// Per-dtype helpers. We avoid heavy template metaprogramming and just provide
// overloaded inline functions; the kernel below selects via `if constexpr`.
// ---------------------------------------------------------------------------
// Type-generic zero — both half and bf16_t in HIP/ROCm have a converting
// constructor from float, but going through __float2half_rn / __float2bfloat16
// is the unambiguously correct path on every ROCm version.
template <typename T>
__forceinline__ __device__ T tzero();
template <>
__forceinline__ __device__ half tzero<half>() {
return __float2half_rn(0.0f);
}
template <>
__forceinline__ __device__ bf16_t tzero<bf16_t>() {
return __float2bfloat16(0.0f);
}
__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
// RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes
// fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator
// staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious
// `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA
// disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in
// the M_COUNT=8 kernel before this change), so we issue the builtin
// explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops)
// per dot22_8_f call vs the half2-accumulator form. With 128 calls per
// K=32 step that's ~384 ops/K-step less issue pressure on the VALU.
//
// Numerical bonus: accumulator stays fp32 throughout the dot. The old
// form accumulated 8 muladds in fp16 (10-bit mantissa) before casting,
// which could lose ~3 bits of precision on borderline magnitudes.
float result = 0.0f;
const half2* a2_ptr = (const half2*)a_ptr;
#pragma unroll
for (int i = 0; i < 4; i++) {
result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
}
return result;
}
__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4],
const bf16_t* a_ptr) {
// RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in
// the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc
// therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised
// fallback (single-element FMAs or fp32 round-trips), which empirically
// runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16
// decode path was paying that tax in full, scaling linearly with M (the
// fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate
// and the kernel becomes memory-bound).
//
// Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs)
// and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA
// count, but each FMA is fast. Bonus: the accumulator is now fp32
// throughout instead of bf16, which is also numerically more accurate
// (no compounding bf16-rounding inside the dot loop).
float result = 0.0f;
#pragma unroll
for (int i = 0; i < 4; i++) {
uint32_t aw, dw;
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
__builtin_memcpy(&dw, &dq[i], sizeof(uint32_t));
// bf16 in low 16 bits → fp32 by left-shifting into the upper half.
// bf16 in high 16 bits → already aligned with fp32's upper half.
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
float a_y = __uint_as_float(aw & 0xFFFF0000u);
float d_x = __uint_as_float((dw & 0xFFFFu) << 16);
float d_y = __uint_as_float(dw & 0xFFFF0000u);
result = __fmaf_rn(d_x, a_x, result);
result = __fmaf_rn(d_y, a_y, result);
}
return result;
}
// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already
// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t
// overload above does for dq (still need to widen A from bf16). Wins more
// at high N: the bf162_t version's per-call widening cost scales with the
// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only
// for A widening (M_COUNT × 4 × 4 widens, half as many).
__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
const bf16_t* a_ptr) {
float result = 0.0f;
#pragma unroll
for (int i = 0; i < 4; i++) {
uint32_t aw;
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
float a_y = __uint_as_float(aw & 0xFFFF0000u);
result = __fmaf_rn(dq[2 * i + 0], a_x, result);
result = __fmaf_rn(dq[2 * i + 1], a_y, result);
}
return result;
}
// ---------------------------------------------------------------------------
// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
// landed on gfx940 / gfx1250 respectively), so this lowers to
// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to
// write 4 output columns per row in a single atomic operation — half the
// atomic instruction count and half the contention vs two 32-bit CAS calls.
//
// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer +
// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue
// cast pass that an fp32-accumulator design would need.
//
// 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4
// (always multiple of 4), and partition_weight_shape[1] is required to be a
// multiple of 8 by can_implement(), so every (m, n) write target is 8-byte
// aligned. Required by global_atomic_cmpswap_b64.
// ---------------------------------------------------------------------------
__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
half2 v23) {
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
unsigned long long old = *addr_u;
while (true) {
union {
unsigned long long u;
half2 h2[2];
} cur, sum;
cur.u = old;
sum.h2[0] = __hadd2(cur.h2[0], v01);
sum.h2[1] = __hadd2(cur.h2[1], v23);
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
if (prev == old) break;
old = prev;
}
}
__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
bf162_t v23) {
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
unsigned long long old = *addr_u;
while (true) {
union {
unsigned long long u;
bf162_t b2[2];
} cur, sum;
cur.u = old;
sum.b2[0] = __hadd2(cur.b2[0], v01);
sum.b2[1] = __hadd2(cur.b2[1], v23);
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
if (prev == old) break;
old = prev;
}
}
// Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8]
// uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with
// offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two
// uint32 words; in practice within one because n & 7 is 0 or 4.
__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
int (&zeros)[4]) {
int qcol = n / 8;
int shift = (n & 0x07) * 4;
uint32_t d = qzeros_row[qcol] >> shift;
zeros[0] = (int)(d & 0xF);
zeros[1] = (int)((d >> 4) & 0xF);
zeros[2] = (int)((d >> 8) & 0xF);
zeros[3] = (int)((d >> 12) & 0xF);
}
template <typename T>
__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
T (&scales)[4]) {
scales[0] = scales_row[n + 0];
scales[1] = scales_row[n + 1];
scales[2] = scales_row[n + 2];
scales[3] = scales_row[n + 3];
}
// ---------------------------------------------------------------------------
// Main kernel.
// ---------------------------------------------------------------------------
template <typename T, int M_COUNT>
__global__ void gemm_q4_kernel_rdna3(
const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
T* __restrict__ c, const int size_m, const int size_n, const int size_k,
const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
const int t = threadIdx.x;
const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
const int offset_m = blockIdx.y * M_COUNT;
const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
const int n = offset_n + t * 4;
// LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M
// row break the natural 256-element/512-byte alignment that would otherwise
// collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k]
// (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132
// 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for
// all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant.
constexpr int LDS_PAD = 8;
__shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD];
// Stage A: each thread loads 1 K element per M row into LDS (with optional
// act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map.
// For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m
// are zero-padded so the dot product contribution is 0 (we then skip the
// atomic write for those rows below).
//
// M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads
// read the SAME 8-element A window per inner step (a_off is uniform across
// the block), so the cache-line broadcast through L1 makes global reads as
// cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128.
static_assert(BLOCK_KN_SIZE == THREADS_X,
"BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)");
// The M=1 fast path (skip LDS) only has a global-read code path for bf16
// (the v_dot2_f32_bf16 branch). The fp16 inner loop still indexes
// block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through
// LDS even at M=1 to avoid reading uninitialized shared memory.
constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same<T, half>::value;
if constexpr (USE_LDS_A) {
if (offset_k + t < end_k) {
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
T av;
if (offset_m + m < size_m) {
const T* a_row = a + (offset_m + m) * size_k;
if (b_q_perm)
av = a_row[b_q_perm[offset_k + t]];
else
av = a_row[offset_k + t];
} else {
av = tzero<T>(); // zero-pad invalid M rows
}
block_a[m][t] = av;
}
}
// Threads beyond the right edge of N have nothing to do. Note: we must NOT
// return before __syncthreads() if any thread in the block participates in
// the LDS load above — but here all THREADS_X (=256) threads always do,
// regardless of whether their `n` is in bounds.
__syncthreads();
} else if (b_q_perm) {
// bf16 M=1 fast path skips LDS, but its global read below is sequential
// and cannot apply act-order. When a permutation is present, stage the
// single A row through LDS (as fp16 / M>1 do) so the read picks it up.
// b_q_perm is block-uniform, so the __syncthreads is non-divergent.
if (offset_k + t < end_k)
block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]];
__syncthreads();
}
if (n >= size_n) return;
// Group bookkeeping. We require size_k % groups == 0 (groupsize divides K).
const int groupsize = size_k / groups;
int group = offset_k / groupsize;
int nextgroup = (group + 1) * groupsize;
// qweight stride: weights are [K/8, N] uint32 with K packed at dim 0.
int qk = offset_k / 8;
const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
// Per-column dequant constants. We hold one set of (z, y) pairs per column.
// fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
// nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant
// produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA
// bypass for the missing v_pk_fma_bf16 on gfx11.
half2 z1z16_h[4][2], y1y16_h[4][2];
float z_b_f[4], y_b_f[4];
auto refresh_group = [&](int g) {
const uint32_t* qz_row = b_qzeros + g * (size_n / 8);
const T* sc_row = b_scales + g * size_n;
int zeros[4];
T scales[4];
load4_zeros(qz_row, n, zeros);
load4_scales<T>(sc_row, n, scales);
if constexpr (std::is_same<T, half>::value) {
#pragma unroll
for (int i = 0; i < 4; ++i) {
prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i],
z1z16_h[i], y1y16_h[i]);
}
} else {
#pragma unroll
for (int i = 0; i < 4; ++i) {
prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i],
z_b_f[i], y_b_f[i]);
}
}
};
refresh_group(group);
float block_c[M_COUNT][4];
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
#pragma unroll
for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
}
// Note on group-transition granularity: we check `k == nextgroup` at the
// start of each outer iteration (which advances K by 32). This is correct
// when group_size >= 32 OR group_size divides 32 evenly (groupsize is one
// of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the
// inner loop would cross a group boundary between j-iterations; we require
// group_size >= 32 here, mirroring exllama's assumption.
//
// Software pipelining: we issue all 4 vectorized weight loads up front
// before any dequant/FMA depends on them. This gives the AMDGPU backend
// freedom to schedule the global_loads early and overlap their latency
// with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16
// VGPRs in flight per thread, plenty of headroom on RDNA3.
int k = offset_k;
while (k < end_k) {
if (k == nextgroup) {
group++;
nextgroup += groupsize;
refresh_group(group);
}
// Prefetch all four j-iterations' weight words. The compiler emits 4
// global_load_b128 instructions back-to-back; the dependent dequant +
// FMA work below hides their latency.
int4 b_w[4];
#pragma unroll
for (int j = 0; j < 4; ++j) {
b_w[j] = *(const int4*)(b_ptr + j * size_n);
}
b_ptr += 4 * size_n;
#pragma unroll
for (int j = 0; j < 4; ++j) {
const int a_off = (k - offset_k) + 8 * j;
if constexpr (std::is_same<T, half>::value) {
half2 dq[4][4];
dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]);
dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]);
dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]);
dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]);
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
block_c[m][0] += dot22_8_f(dq[0], a_ptr);
block_c[m][1] += dot22_8_f(dq[1], a_ptr);
block_c[m][2] += dot22_8_f(dq[2], a_ptr);
block_c[m][3] += dot22_8_f(dq[3], a_ptr);
}
} else if constexpr (M_COUNT == 1) {
// bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of
// Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's
// InstCombine cannot fold the bf16→fp32 widening (LLVM #76000):
// * activations and magic-value weights share a fp32-aliased
// union (bytes written as uint32, read as bf16x2_t for the
// dot — pointer-cast opacity defeats the fold)
// * sum_a computed via a *second* v_dot2 with bf162(1,1) as the
// second operand, avoiding any explicit bf16→fp32 widen of A
// * bias correction y_b_f * partial + z_b_f * sum_a, identical
// to the previous fp32-FMA-chain path
//
// Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the
// previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so
// the substitution is ~2× cheaper for the inner accumulator.
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128)
constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0)
union pack4 {
float f[4];
uint32_t u[4];
};
uint32_t w[4];
__builtin_memcpy(w, &b_w[j], sizeof(int4));
// Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a
// fp32-aliased union. Storing as uint32 keeps the IR-level type
// opaque so the inner v_dot2 cannot be folded to fp32 widening.
//
// A is read direct from global (no LDS staging — see USE_LDS_A above),
// except under act-order, where it comes from the permuted LDS copy.
pack4 a_pack;
{
const uint32_t* a_words =
b_q_perm
? reinterpret_cast<const uint32_t*>(&block_a[0][a_off])
: reinterpret_cast<const uint32_t*>(a + offset_k + a_off);
a_pack.u[0] = a_words[0];
a_pack.u[1] = a_words[1];
a_pack.u[2] = a_words[2];
a_pack.u[3] = a_words[3];
}
// sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as
// the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi.
// No fp32 widening of activations: the bytes go straight from LDS
// through v_dot2 into the fp32 accumulator.
float sum_a = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
sum_a = __builtin_amdgcn_fdot2_f32_bf16(
*((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
sum_a, /*clamp=*/false);
}
// unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs
// recycled across cols), avoiding straight-line expansion that
// would inflate live-range to 32 VGPRs.
#pragma unroll 1
for (int col = 0; col < 4; ++col) {
// Build dequant magic values bf16(128 + nibble) directly into a
// fp32-aliased union via uint32 stores. No fp32 in the data flow
// until v_dot2 consumes the bytes.
pack4 q_pack;
const uint32_t qa = w[col];
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
// partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16.
float partial = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
partial = __builtin_amdgcn_fdot2_f32_bf16(
*((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
partial, /*clamp=*/false);
}
// block_c += y_b_f * partial + z_b_f * sum_a
// y_b_f = scale, z_b_f = -(128+zero)*scale
// partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a
// and scaling yields scale · (nibble - zero) · a as required.
block_c[0][col] =
__fmaf_rn(y_b_f[col], partial,
__fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
}
} else {
// bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as
// the M=1 branch: activations + magic-value weights stored in
// fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with
// pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2
// with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a)
// applied after the dot. Magic values built once per col and reused
// across all M rows — amortizes dequant cost across M_COUNT.
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128)
constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0)
union pack4 {
float f[4];
uint32_t u[4];
};
uint32_t w[4];
__builtin_memcpy(w, &b_w[j], sizeof(int4));
// Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4
// unions. Stored as uint32 to keep IR-level types opaque (defeats
// InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3
// budget.
pack4 a_pack[M_COUNT];
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
const uint32_t* a_words =
reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
a_pack[m].u[0] = a_words[0];
a_pack[m].u[1] = a_words[1];
a_pack[m].u[2] = a_words[2];
a_pack[m].u[3] = a_words[3];
}
// sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen.
float sum_a[M_COUNT];
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
float s = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
*((const bf16x2_t*)&BF16_ONES),
s, /*clamp=*/false);
}
sum_a[m] = s;
}
// Per col: build magic-value pack, dot against all M activations.
// unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled)
// — same register-pressure trick as the previous fp32 path.
#pragma unroll 1
for (int col = 0; col < 4; ++col) {
pack4 q_pack;
const uint32_t qa = w[col];
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
float partial = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
partial = __builtin_amdgcn_fdot2_f32_bf16(
*((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
partial, /*clamp=*/false);
}
// block_c += y_b_f * partial + z_b_f * sum_a (same correction as
// M=1)
block_c[m][col] =
__fmaf_rn(y_b_f[col], partial,
__fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
}
}
}
}
k += 32; // 4 weight words * 8 nibbles = 32 K elements
}
// Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all
// four lanes in a single 64-bit CAS write directly to the T-typed output
// (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with
// a single b64 op we halve the atomic instruction count vs two b32 CAS
// calls, AND save the FP32 buffer + memset + cast pass entirely.
#pragma unroll
for (int m = 0; m < M_COUNT; ++m) {
if (offset_m + m >= size_m) continue; // skip padding rows past size_m
T* out = c + (offset_m + m) * size_n + n;
if constexpr (std::is_same<T, half>::value) {
half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
__float2half_rn(block_c[m][1]));
half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
__float2half_rn(block_c[m][3]));
atomic_add_pk4_f16(out, r01, r23);
} else {
bf162_t r01;
r01.x = __float2bfloat16(block_c[m][0]);
r01.y = __float2bfloat16(block_c[m][1]);
bf162_t r23;
r23.x = __float2bfloat16(block_c[m][2]);
r23.y = __float2bfloat16(block_c[m][3]);
atomic_add_pk4_bf16(out, r01, r23);
}
}
}
#else // non-RDNA3 device pass: empty __global__ for symbol parity.
template <typename T, int M_COUNT>
__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
const T*, T*, const int, const int,
const int, const int, const int,
const int*) {}
#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
// ---------------------------------------------------------------------------
// Launcher.
// ---------------------------------------------------------------------------
template <typename T, int M_COUNT>
void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight,
const uint32_t* b_qzeros, const T* b_scales,
const int* b_q_perm, T* c, int size_m,
int size_n, int size_k, int groups,
int zero_offset, cudaStream_t stream) {
dim3 block(THREADS_X);
dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
(size_m + M_COUNT - 1) / M_COUNT,
(size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
zero_offset, b_q_perm);
}
// Dispatch to the largest M_COUNT template that doesn't waste more than
// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is
// the right tool, not bigger M_COUNT in the scalar dot-product path.
//
// Tile-waste table:
// M=1 -> M_COUNT=1 (no waste)
// M=2,3 -> M_COUNT=2 (M=3 wastes 1/2 of last tile)
// M=4-7 -> M_COUNT=4 (worst case M=5: wastes 3/4 of last tile)
// M=8-15-> M_COUNT=8 (worst case M=9: wastes 7/8 of last tile)
// "Wasted" rows are zero-padded in LDS and skip the atomic write, so they
// only burn instructions on the last block, never affect correctness.
template <typename T>
void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
const uint32_t* b_qzeros, const T* b_scales,
const int* b_q_perm, T* c, int size_m, int size_n,
int size_k, int groups, bool use_v2_format,
cudaStream_t stream) {
const int zero_offset = use_v2_format ? 0 : 1;
if (size_m == 1) {
launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
c, size_m, size_n, size_k, groups,
zero_offset, stream);
} else if (size_m <= 3) {
launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
c, size_m, size_n, size_k, groups,
zero_offset, stream);
} else if (size_m <= 7) {
launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
c, size_m, size_n, size_k, groups,
zero_offset, stream);
} else {
// M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the
// WMMA path, but if it falls through we still produce correct output —
// just leaving 3-5× of throughput on the table for prefill workloads.
launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
c, size_m, size_n, size_k, groups,
zero_offset, stream);
}
}
} // namespace gptq_rdna3
} // namespace vllm
// ---------------------------------------------------------------------------
// Public entry point.
// ---------------------------------------------------------------------------
//
// Inputs:
// a [M, K] half or bfloat16
// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle)
// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros)
// b_scales [groups, N] half or bfloat16
// b_g_idx [K] or empty int32 (act-order permutation; empty=identity)
// use_v2_format bool (true = GPTQv2, no +1 zero offset)
//
// Output:
// c [M, N] same dtype as a
torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros,
torch::Tensor b_scales,
torch::Tensor b_g_idx, bool use_v2_format);
torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros, torch::Tensor b_scales,
torch::Tensor b_g_idx, bool use_v2_format) {
if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 &&
b_q_weight.size(1) % 16 == 0 &&
((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) ||
(a.scalar_type() == torch::kHalf && a.size(0) >= 64))) {
return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx,
use_v2_format);
}
TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor");
TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor");
TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]");
TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]");
TORCH_CHECK(
a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
"a must be half or bfloat16");
TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
"b_scales dtype must match a");
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
auto stream = at::cuda::getCurrentCUDAStream();
int size_m = (int)a.size(0);
int size_k = (int)a.size(1);
int size_n = (int)b_q_weight.size(1);
int groups = (int)b_qzeros.size(0);
TORCH_CHECK(b_q_weight.size(0) * 8 == size_k,
"b_q_weight first dim must be K/8");
TORCH_CHECK(b_scales.size(0) == groups,
"b_scales must have same group count as qzeros");
TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N");
TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)");
auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
at::Tensor c = torch::zeros({size_m, size_n}, opts);
const int* g_idx_ptr = nullptr;
if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
"b_g_idx must be int32");
g_idx_ptr = (const int*)b_g_idx.data_ptr();
}
if (a.scalar_type() == torch::kHalf) {
vllm::gptq_rdna3::launch_gemm_q4<half>(
(const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
(const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
use_v2_format, stream);
} else {
vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>(
(const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
(const uint32_t*)b_q_weight.data_ptr(),
(const uint32_t*)b_qzeros.data_ptr(),
(const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
(vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups,
use_v2_format, stream);
}
return c;
}
csrc/rocm/q_gemm_rdna3_wmma.cu
+2165
View File
File diff suppressed because it is too large Load Diff
csrc/rocm/qdq_4_rdna3.cuh
+239
View File
@@ -0,0 +1,239 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on
// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses
// the classic exllamav2 bit-trick:
//
// (qa & 0x000F000F) | 0x64006400 -> half2(1024+q_lo, 1024+q_hi)
// (qa & 0x00F000F0) | 0x64006400 -> half2(1024+q_lo*16, 1024+q_hi*16)
//
// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs
// works in fp16 because the mantissa (10 bits) is wide enough to hold a value
// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper
// nibble into bits [7:4] would spill into the exponent. To avoid that, the
// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a
// single right-shift before the OR with 0x43004300 (= bf162(128, 128)).
#ifndef _qdq_4_rdna3_cuh
#define _qdq_4_rdna3_cuh
#include <cstdint>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
namespace vllm {
namespace gptq_rdna3 {
using bf16_t = __hip_bfloat16;
using bf162_t = __hip_bfloat162;
// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]:
// in: q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0] (LSB first)
// out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0] (even/odd interleaved)
//
// After shuffle, q[2k] sits at bits [4k : 4k+3] (lower 16)
// q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16)
// so a single mask 0x000F000F selects the matching even/odd pair, ready to
// bitcast to half2 / bfloat162 after OR-ing with the magic constant.
__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) {
uint32_t qa = q[0];
uint32_t qb = 0;
#pragma unroll
for (int i = 0; i < 4; i++) {
uint32_t qa0 = qa & 0x0F;
uint32_t qa1 = (qa & 0xF0) >> 4;
qa >>= 8;
qb |= (qa1 << (i * 4 + 16));
qb |= (qa0 << (i * 4));
}
q[0] = qb;
}
// ---------------------------------------------------------------------------
// fp16 path
// ---------------------------------------------------------------------------
// Precompute scale-baked constants for a single zero/scale pair.
// z1z16[0] = scale * (-1024 - zero) (used for "low" pairs)
// z1z16[1] = scale * (-64 - zero) (used for "high" pairs)
// y1y16[0] = scale * 1 (low pairs are q + 1024)
// y1y16[1] = scale * (1/16) (high pairs are q*16 + 1024)
__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale,
half2 (&z1z16)[2],
half2 (&y1y16)[2]) {
// half(-1024 - zero) via the exllamav2 bit-trick:
// half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it.
union {
uint16_t u;
half h;
} z1u;
z1u.u = (uint16_t)(0xE400 | zero);
half z1 = z1u.h;
half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero));
half2 scale2 = __half2half2(scale);
z1z16[0] = __hmul2(scale2, __half2half2(z1));
z1z16[1] = __hmul2(scale2, __half2half2(z16));
half y1 = __float2half_rn(1.0f);
half y16 = __float2half_rn(1.0f / 16.0f);
y1y16[0] = __hmul2(scale2, __half2half2(y1));
y1y16[1] = __hmul2(scale2, __half2half2(y16));
}
// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs:
// dq[0] = (q[0], q[1]) * scale - zero*scale
// dq[1] = (q[2], q[3]) * scale - zero*scale
// dq[2] = (q[4], q[5]) * scale - zero*scale
// dq[3] = (q[6], q[7]) * scale - zero*scale
__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
half2 (&z1z16)[2],
half2 (&y1y16)[2]) {
const uint32_t c0 = 0x64006400;
union {
uint32_t u;
half2 h2;
} q0, q1, q2, q3;
q0.u = (qa & 0x000F000F) | c0; // half2(q[0]+1024, q[1]+1024)
q1.u = (qa & 0x00F000F0) | c0; // half2(q[2]*16+1024, q[3]*16+1024)
uint32_t qa_hi = qa >> 8;
q2.u = (qa_hi & 0x000F000F) | c0; // half2(q[4]+1024, q[5]+1024)
q3.u = (qa_hi & 0x00F000F0) | c0; // half2(q[6]*16+1024, q[7]*16+1024)
dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]);
dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]);
dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]);
dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]);
}
// ---------------------------------------------------------------------------
// bf16 path
// ---------------------------------------------------------------------------
// Bit-trick magic for bf16:
// bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0).
// For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because
// bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So
// ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi).
//
// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble
// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16]
// via a single 4/8/12-bit right-shift before the OR. That costs one extra
// shift per pair vs fp16, but keeps the FMA structure identical.
__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero,
bf16_t scale,
bf162_t& z_prep,
bf162_t& y_prep) {
// z = scale * -(128 + zero); y = scale.
float scale_f = __bfloat162float(scale);
float zf = -(128.0f + (float)zero) * scale_f;
bf16_t zb = __float2bfloat16(zf);
z_prep = __bfloat162bfloat162(zb);
y_prep = __bfloat162bfloat162(scale);
}
__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa,
bf162_t (&dq)[4],
bf162_t z_prep,
bf162_t y_prep) {
const uint32_t c0 = 0x43004300;
union {
uint32_t u;
bf162_t b2;
} q0, q1, q2, q3;
q0.u = ((qa >> 0) & 0x000F000F) | c0; // bf162(128+q[0], 128+q[1])
q1.u = ((qa >> 4) & 0x000F000F) | c0; // bf162(128+q[2], 128+q[3])
q2.u = ((qa >> 8) & 0x000F000F) | c0; // bf162(128+q[4], 128+q[5])
q3.u = ((qa >> 12) & 0x000F000F) | c0; // bf162(128+q[6], 128+q[7])
// dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale
dq[0] = __hfma2(q0.b2, y_prep, z_prep);
dq[1] = __hfma2(q1.b2, y_prep, z_prep);
dq[2] = __hfma2(q2.b2, y_prep, z_prep);
dq[3] = __hfma2(q3.b2, y_prep, z_prep);
}
// ---------------------------------------------------------------------------
// bf16-input → fp32-output dequant (RDNA3 scalar path).
//
// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow
// fallback. Rather than computing dq in bf16 and widening at FMA time in
// the dot product, we widen to fp32 here once (a free left-shift by 16) and
// emit the (q - zero) * scale FMA directly in fp32. This:
// * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32.
// * Eliminates 4× bf16→fp32 widens that the dot product would do.
// * Keeps the dot product accumulator in fp32 without a roundtrip.
//
// Output: fp32 dq[8], one element per K position (consumed by the
// fp32-overload of dot22_8_f in q_gemm_rdna3.cu).
__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
bf16_t scale,
float& z_prep,
float& y_prep) {
float scale_f = __bfloat162float(scale);
z_prep = -(128.0f + (float)zero) * scale_f;
y_prep = scale_f;
}
// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32
// values 128+nibble, without folding scale/zero. The caller folds scale/zb
// into the accumulator outside the inner loop using a precomputed sum_a,
// which saves ~27% of the FMA count vs the per-col-dequant approach above
// (only beneficial at M_COUNT=1; break-even at M_COUNT=2).
//
// Cost: 0 FMAs (pure bit-trick + as_float reinterprets).
__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa,
float (&q_f32)[8]) {
const uint32_t c0 = 0x43004300;
const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16);
q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u);
q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16);
q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u);
q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16);
q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u);
q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16);
q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u);
}
__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa,
float (&dq)[8],
float z_prep,
float y_prep) {
const uint32_t c0 = 0x43004300;
const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
// bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16
// (just zero-extends the mantissa from 7 to 23 bits; exponent preserved).
const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
// dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale
dq[0] = __fmaf_rn(q0x, y_prep, z_prep);
dq[1] = __fmaf_rn(q0y, y_prep, z_prep);
dq[2] = __fmaf_rn(q1x, y_prep, z_prep);
dq[3] = __fmaf_rn(q1y, y_prep, z_prep);
dq[4] = __fmaf_rn(q2x, y_prep, z_prep);
dq[5] = __fmaf_rn(q2y, y_prep, z_prep);
dq[6] = __fmaf_rn(q3x, y_prep, z_prep);
dq[7] = __fmaf_rn(q3y, y_prep, z_prep);
}
} // namespace gptq_rdna3
} // namespace vllm
#endif // _qdq_4_rdna3_cuh
csrc/rocm/torch_bindings.cpp
+13
View File
@@ -39,6 +39,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
" Tensor scale_b, int CuCount) -> ()");
rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
#ifdef VLLM_ROCM_GFX1100
// W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
rocm_ops.def(
"gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
"Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3);
rocm_ops.def(
"gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
"Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
#endif
// Custom attention op
// Compute the attention between an input query and the cached
// keys/values using PagedAttention.
csrc/torch_bindings.cpp
+4 -138
View File
@@ -1,4 +1,7 @@
#include "cache.h"
// Provides torch::Tensor for ops.h (previously included transitively via
// cache.h, which is no longer included here after cache ops moved to
// _C_stable_libtorch).
#include <torch/all.h>
#include "cuda_utils.h"
#include "ops.h"
#include "core/registration.h"
@@ -33,35 +36,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
&get_cuda_view_from_cpu_tensor);
// Attention ops
// Compute the attention between an input query and the cached
// keys/values using PagedAttention.
ops.def(
"paged_attention_v1("
" Tensor! out, Tensor query, Tensor key_cache,"
" Tensor value_cache, int num_kv_heads, float scale,"
" Tensor block_tables, Tensor seq_lens, int block_size,"
" int max_seq_len, Tensor? alibi_slopes,"
" str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
" int tp_rank, int blocksparse_local_blocks,"
" int blocksparse_vert_stride, int blocksparse_block_size,"
" int blocksparse_head_sliding_step) -> ()");
ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
// PagedAttention V2.
ops.def(
"paged_attention_v2("
" Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
" Tensor! tmp_out, Tensor query, Tensor key_cache,"
" Tensor value_cache, int num_kv_heads, float scale,"
" Tensor block_tables, Tensor seq_lens, int block_size,"
" int max_seq_len, Tensor? alibi_slopes,"
" str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
" int tp_rank, int blocksparse_local_blocks,"
" int blocksparse_vert_stride, int blocksparse_block_size,"
" int blocksparse_head_sliding_step) -> ()");
ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -217,114 +191,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#endif
}
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
// Cache ops
// Swap in (out) the cache blocks from src to dst.
cache_ops.def(
"swap_blocks(Tensor src, Tensor! dst,"
" int block_size_in_bytes, Tensor block_mapping) -> ()");
cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
// Batch swap: submit all block copies in a single driver call.
cache_ops.def(
"swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
" Tensor sizes,"
" bool is_src_access_order_any=False) -> ()");
cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
// Reshape the key and value tensors and cache them.
cache_ops.def(
"reshape_and_cache(Tensor key, Tensor value,"
" Tensor! key_cache, Tensor! value_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale) -> ()");
cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache);
// Reshape the key and value tensors and cache them.
cache_ops.def(
"reshape_and_cache_flash(Tensor key, Tensor value,"
" Tensor! key_cache,"
" Tensor! value_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale) -> ()");
cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
&reshape_and_cache_flash);
// Concat kv_c and k_pe and cache them.
cache_ops.def(
"concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
" Tensor! kv_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor scale) -> ()");
cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
// Rotate Q and K, then write to kv cache for MLA
cache_ops.def(
"concat_and_cache_mla_rope_fused("
" Tensor positions,"
" Tensor! q_pe,"
" Tensor! k_pe,"
" Tensor kv_c,"
" Tensor cos_sin_cache,"
" bool is_neox,"
" Tensor slot_mapping,"
" Tensor! kv_cache,"
" str kv_cache_dtype,"
" Tensor kv_cache_scale) -> ()");
cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
&concat_and_cache_mla_rope_fused);
// Convert the key and value cache to fp8 data type.
cache_ops.def(
"convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
"str kv_cache_dtype) -> ()");
cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
// Gather cache blocks from src_cache to dst, dequantizing from
// src_cache's dtype to dst's dtype if necessary.
cache_ops.def(
"gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
" Tensor block_table, Tensor cu_seq_lens, "
" Tensor token_to_seq, "
" int num_tokens, "
" str kv_cache_dtype, "
" Tensor scale, Tensor? seq_starts) -> ()");
cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
&gather_and_maybe_dequant_cache);
cache_ops.def(
"cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
"Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
cache_ops.def(
"cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
"Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
"batch_size) -> ()");
cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA,
&cp_gather_and_upconvert_fp8_kv_cache);
cache_ops.def(
"indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
"slot_mapping, "
"int quant_block_size, str kv_cache_dtype) -> ()");
cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
&indexer_k_quant_and_cache);
cache_ops.def(
"concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
cache_ops.def(
"cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
"dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA,
&cp_gather_indexer_k_quant_cache);
}
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
// Cuda utils
docker/Dockerfile
+2 -2
View File
@@ -757,10 +757,10 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.11.post2
ARG FLASHINFER_VERSION=0.6.12
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
# ============================================================
# OPENAI API SERVER DEPENDENCIES
docker/Dockerfile.cpu
+25 -1
View File
@@ -165,6 +165,23 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
######################### TRITON-CPU BUILD IMAGE #########################
FROM base AS vllm-triton-cpu-build
WORKDIR /vllm-workspace
RUN mkdir dist
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
git clone --recurse-submodules "https://github.com/triton-lang/triton-cpu.git"; \
cd triton-cpu; \
git checkout "270e696d"; \
uv build --wheel --out-dir=../dist; \
fi
######################### TEST DEPS #########################
FROM base AS vllm-test-deps
@@ -257,7 +274,14 @@ WORKDIR /vllm-workspace
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/root/.cache/ccache \
--mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
uv pip install "$(realpath dist/*.whl)[audio]"
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/root/.cache/ccache \
--mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
uv pip install "$(realpath dist/*.whl)"; \
fi
# Add labels to document build configuration
LABEL org.opencontainers.image.title="vLLM CPU"

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