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-support-rms-norm-uncontiguous

Browse Source 
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
This commit is contained in:
Wentao Ye
2026-07-28 14:31:16 -04:00
committed by GitHub
parent 4ae4262b17 6c7e679f04
commit 59c62cd2b4
661 changed files with 40037 additions and 10421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/check-torch-abi.py
+102
View File
@@ -0,0 +1,102 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Audit vLLM compiled libraries for PyTorch stable ABI compliance."""
import fnmatch
import sys
from pathlib import Path
from torch_abi_audit import inspect_package
from torch_abi_audit.report import ExtensionReport, PackageReport
# Temporary allowlist of extensions not yet on the stable ABI.
# Shrink and remove over time.
ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
"vllm_flash_attn/_vllm_fa2_C.abi3.so",
"vllm_flash_attn/_vllm_fa3_C.abi3.so",
"third_party/deep_gemm/_C*.so",
)
def _relative_path(lib: ExtensionReport, package_root: Path) -> str:
try:
return lib.path.relative_to(package_root).as_posix()
except ValueError:
return lib.path.name
def _is_torch_unstable(lib: ExtensionReport) -> bool:
return lib.error is None and lib.torch.uses_torch and not lib.torch.stable
def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool:
return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns)
def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]:
return (*report.extensions, *report.bundled_libs)
def _collect_unstable(report: PackageReport) -> list[str]:
return sorted(
_relative_path(lib, report.root)
for lib in _iter_libs(report)
if _is_torch_unstable(lib)
)
def _find_stale_allowlist_entries(
report: PackageReport, patterns: tuple[str, ...]
) -> list[str]:
"""Allowlist patterns that match a built library which is no longer unstable."""
stale: list[str] = []
for pattern in patterns:
for lib in _iter_libs(report):
if lib.error is not None:
continue
if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern):
continue
if not _is_torch_unstable(lib):
stale.append(pattern)
break
return stale
def check_torch_abi(
package: str = "vllm",
patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES,
) -> int:
report = inspect_package(package)
if report.error:
print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr)
return 2
unstable = _collect_unstable(report)
unexpected = [
rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns)
]
stale = _find_stale_allowlist_entries(report, patterns)
if unexpected or stale:
if unexpected:
print(
"Not allowed: torch-unstable libraries outside "
f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}",
file=sys.stderr,
)
if stale:
print(
"Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: "
f"{', '.join(stale)}",
file=sys.stderr,
)
return 1
print("Torch stable ABI check passed.")
return 0
if __name__ == "__main__":
print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance")
sys.exit(check_torch_abi())
.buildkite/ci_config.yaml
+1
View File
@@ -14,6 +14,7 @@ run_all_patterns:
- "setup.py"
- "csrc/"
- "cmake/"
- ".buildkite/check-torch-abi.py"
run_all_exclude_patterns:
- "docker/Dockerfile."
- "csrc/cpu/"
.buildkite/intel_jobs/benchmarks_intel.yaml
+26
View File
@@ -0,0 +1,26 @@
group: Benchmarks
depends_on:
- image-build-xpu
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 40
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/benchmarks/
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s benchmarks/'
.buildkite/intel_jobs/engine_intel.yaml
+76
View File
@@ -2,6 +2,44 @@ group: Engine Intel
depends_on:
- image-build-xpu
steps:
- label: Engine
key: engine
timeout_in_minutes: 40
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/compilation/
- vllm/config/
- vllm/engine/
- vllm/entrypoints/logger.py
- vllm/envs.py
- vllm/logger.py
- vllm/logging_utils/
- vllm/platforms/
- vllm/sequence.py
- vllm/triton_utils/
- vllm/utils/
- tests/engine
- tests/test_sequence
- tests/test_config
- tests/test_logger
- tests/test_vllm_port
- tests/test_jit_monitor.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py'
- label: Engine (1 GPU)
timeout_in_minutes: 30
device: intel_gpu
@@ -23,3 +61,41 @@ steps:
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
- label: V1 e2e (2 GPUs)
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/compilation/
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/envs.py
- vllm/forward_context.py
- vllm/inputs/
- vllm/logger.py
- vllm/logging_utils/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/triton_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/e2e/spec_decode
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"'
.buildkite/intel_jobs/misc_intel.yaml
+29 -4
View File
@@ -125,13 +125,13 @@ steps:
pytest -v -s v1/kv_offload &&
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
- label: NixlConnector PD accuracy (2 GPUs)
- label: NixlConnector PD accuracy (4 GPUs)
timeout_in_minutes: 60
num_devices: 2
num_devices: 4
device: intel_gpu
agent_tags:
label: production
gpu: 2+
gpu: 4+
mem: 16+
no_plugin: true
working_dir: "."
@@ -148,7 +148,10 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
- label: Regression
key: regression
@@ -259,3 +262,25 @@ steps:
pytest -v -s detokenizer &&
pytest -v -s -m "not cpu_test" ./multimodal &&
pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
- label: Fusion Unit Tests
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/compilation/
- tests/compile/passes/test_qk_norm_rope_fusion.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
.buildkite/intel_jobs/model_executor_intel.yaml
+33
View File
@@ -0,0 +1,33 @@
group: Model Executor Intel
depends_on:
- image-build-xpu
steps:
- label: Model Executor (Intel)
key: model-executor-intel
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'apt-get update && apt-get install -y curl libsodium23 &&
pip3 install tensorizer==2.10.1 &&
pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
export PYTHONFAULTHANDLER=1 &&
cd tests &&
pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"'
.buildkite/intel_jobs/model_runner_v2_intel.yaml
+55 -1
View File
@@ -8,7 +8,7 @@ steps:
agent_tags:
label: production
gpu: 2+
mem: 16+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -28,7 +28,9 @@ steps:
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
cd tests &&
pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
pytest -v -s v1/e2e/general/test_context_length.py &&
ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" &&
pytest -v -s v1/e2e/general/test_min_tokens.py'
- label: Model Runner V2 Examples (Intel)
@@ -60,3 +62,55 @@ steps:
python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
python3 generate/multimodal/vision_language_offline.py --seed 0 &&
python3 features/automatic_prefix_caching/prefix_caching_offline.py'
- label: Model Runner V2 Distributed (2 GPUs)
timeout_in_minutes: 50
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/basic_correctness/test_basic_correctness.py
- tests/v1/distributed/test_async_llm_dp.py
- tests/v1/distributed/test_eagle_dp.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
cd tests &&
TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"'
- label: Model Runner V2 Spec Decode
timeout_in_minutes: 50
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'export VLLM_USE_V2_MODEL_RUNNER=1 &&
cd tests &&
pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py'
.buildkite/intel_jobs/samplers_intel.yaml
+29
View File
@@ -0,0 +1,29 @@
group: Samplers Intel
depends_on:
- image-build-xpu
steps:
- label: Samplers Test (FlashInfer)
key: samplers-test-flashinfer-intel
timeout_in_minutes: 40
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/model_executor/layers
- vllm/sampling_metadata.py
- tests/samplers
- tests/conftest.py
- vllm/entrypoints/generate/beam_search
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers'
.buildkite/intel_jobs/test-intel.yaml
+3 -2
View File
@@ -145,7 +145,8 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_auto_round.py'
pytest -v -s quantization/test_auto_round.py &&
pytest -v -s quantization/test_online.py'
- label: "XPU compressed tensors FP8 test"
depends_on:
- image-build-xpu
@@ -168,4 +169,4 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
.buildkite/scripts/build-macos-wheel.sh
+3
View File
@@ -7,6 +7,9 @@
set -euo pipefail
# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
git fetch --tags --force origin
# The Rust frontend build needs protoc.
if ! command -v protoc >/dev/null 2>&1; then
brew install protobuf
.buildkite/scripts/hardware_ci/run-amd-test.sh
+25 -6
View File
@@ -387,6 +387,7 @@ initialize_native_environment() {
local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
local job_id_suffix=""
local native_root=""
local hf_fstype=""
local hf_mount=""
if [[ "$(id -u)" -ne 0 ]]; then
@@ -400,16 +401,19 @@ initialize_native_environment() {
native_root="/tmp/vllm-native-${job_id}"
TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
VLLM_RPC_BASE_PATH="/tmp"
: "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}"
: "${TRITON_CACHE_DIR:=${native_root}/cache/triton}"
: "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}"
: "${XDG_CACHE_HOME:=${native_root}/cache/xdg}"
TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
TRITON_CACHE_DIR="${native_root}/cache/triton"
VLLM_CACHE_ROOT="${native_root}/cache/vllm"
XDG_CACHE_HOME="${native_root}/cache/xdg"
: "${HF_HOME:=/home/buildkite-agent/huggingface}"
# datasets uses POSIX locks that are unsupported by the shared HF NFS cache.
# Keep processed datasets job-local while retaining the persistent Hub cache.
HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets"
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
: "${HF_HUB_ETAG_TIMEOUT:=60}"
export TMPDIR VLLM_RPC_BASE_PATH
export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
export PYTORCH_ROCM_ARCH=""
mkdir -p "${TMPDIR}" \
@@ -417,7 +421,10 @@ initialize_native_environment() {
"${TRITON_CACHE_DIR}" \
"${VLLM_CACHE_ROOT}" \
"${XDG_CACHE_HOME}" \
"${HF_HOME}" || return 1
"${HF_HOME}" \
"${HF_DATASETS_CACHE}" || return 1
echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
if ! command -v findmnt >/dev/null 2>&1; then
@@ -430,6 +437,18 @@ initialize_native_environment() {
return 1
fi
fi
if command -v findmnt >/dev/null 2>&1; then
hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true)
fi
if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then
# Keep hf-xet state local and avoid vectored writes on shared NFS.
export HF_XET_CACHE="${native_root}/cache/hf-xet"
export HF_XET_HIGH_PERFORMANCE=0
export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0
mkdir -p "${HF_XET_CACHE}" || return 1
echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}"
fi
}
run_native_preflight() {
.buildkite/test-amd.yaml
+66 -41
View File
@@ -81,10 +81,8 @@
# the above test.) Also run if model initialization test file is modified. #
# * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model #
# source is modified, or when specified test files are modified. #
# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to #
# run plamo2 model in vLLM. #
# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) #
# and to run plamo2 model in vLLM. #
# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). #
# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). #
# * [Multi-Modal Models (Standard) 1-4]: #
# - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. #
# * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. #
@@ -171,20 +169,6 @@ steps:
- pip install helion==1.1.0
- pytest -v -s kernels/helion/
- label: Kernels Mamba Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/mamba/
- tests/kernels/mamba
- vllm/model_executor/layers/mamba/ops
- vllm/platforms/rocm.py
commands:
- pytest -v -s kernels/mamba
#------------------------------------------------------ mi250 · models / basic -------------------------------------------------------#
- label: Basic Models Test (Other CPU) # TBD
@@ -366,22 +350,6 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: V1 attention (H100-MI250) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/config/attention.py
- vllm/model_executor/layers/attention
- vllm/v1/attention
- tests/v1/attention
- vllm/_aiter_ops.py
- vllm/envs.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/attention
- label: V1 others (CPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -1579,11 +1547,12 @@ steps:
commands:
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: Kernels Core Operation Test # TBD
- label: Kernels Core Operation Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 3
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/
@@ -1594,7 +1563,7 @@ steps:
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: Kernels KDA Test # TBD
timeout_in_minutes: 180
@@ -1612,6 +1581,21 @@ steps:
commands:
- pytest -v -s kernels/test_kda.py
- label: Kernels Mamba Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/mamba/
- tests/kernels/mamba
- vllm/model_executor/layers/mamba/ops
- vllm/platforms/rocm.py
commands:
- pytest -v -s kernels/mamba
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2163,7 +2147,7 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 4
parallelism: 8
optional: true
working_dir: "/vllm-workspace/"
source_file_dependencies:
@@ -2696,11 +2680,12 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
- label: V1 attention (H100-MI300) # TBD
- label: V1 attention (H100-MI300) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
parallelism: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -2712,7 +2697,7 @@ steps:
- vllm/envs.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/attention
- pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: V1 Core + KV + Metrics # TBD
timeout_in_minutes: 180
@@ -3048,6 +3033,7 @@ steps:
- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/"
@@ -3064,6 +3050,7 @@ steps:
- label: Distributed Tests (2xH100-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3108,6 +3095,7 @@ steps:
- label: Entrypoints Integration (API Server) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
fast_check: true
@@ -3125,6 +3113,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
optional: true
@@ -3140,6 +3129,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
optional: true
@@ -3156,6 +3146,7 @@ steps:
- label: Entrypoints Integration (API Server Generate) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
optional: true
@@ -3176,6 +3167,7 @@ steps:
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
working_dir: "/vllm-workspace/tests"
@@ -3189,6 +3181,7 @@ steps:
- label: Entrypoints Integration (Multimodal)
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
working_dir: "/vllm-workspace/tests"
@@ -3202,6 +3195,7 @@ steps:
- label: Entrypoints Integration (Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
fast_check: true
working_dir: "/vllm-workspace/tests"
@@ -3217,6 +3211,7 @@ steps:
- label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3239,6 +3234,7 @@ steps:
- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3261,6 +3257,7 @@ steps:
- label: LM Eval Small Models (2xB200-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
optional: true
@@ -3280,6 +3277,7 @@ steps:
- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace"
@@ -3300,6 +3298,7 @@ steps:
- label: LM Eval Large Models (4xH100-4xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_4
num_gpus: 4
optional: true
@@ -3322,6 +3321,7 @@ steps:
- label: Examples # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
@@ -3357,6 +3357,7 @@ steps:
- label: Kernels (B200-MI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/"
source_file_dependencies:
@@ -3382,6 +3383,7 @@ steps:
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 2
working_dir: "/vllm-workspace/tests"
@@ -3399,6 +3401,7 @@ steps:
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 5
working_dir: "/vllm-workspace/tests"
@@ -3419,6 +3422,7 @@ steps:
- label: Kernels Quantization Test %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 2
working_dir: "/vllm-workspace/tests"
@@ -3436,6 +3440,7 @@ steps:
- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
@@ -3455,6 +3460,7 @@ steps:
- label: Language Models Test (Extended Generation) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3468,6 +3474,7 @@ steps:
- label: Language Models Test (Extended Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3480,6 +3487,7 @@ steps:
- label: Language Models Test (PPL) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3508,6 +3516,7 @@ steps:
- label: Language Models Tests (Standard) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3522,6 +3531,7 @@ steps:
- label: Multi-Modal Models (Extended Generation 1) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3536,6 +3546,7 @@ steps:
- label: Multi-Modal Models (Extended Generation 3) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3548,6 +3559,7 @@ steps:
- label: Multi-Modal Models (Extended Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3560,6 +3572,7 @@ steps:
- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3573,6 +3586,7 @@ steps:
- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3589,6 +3603,7 @@ steps:
- label: Quantized Models Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3605,6 +3620,7 @@ steps:
- label: Quantization # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3621,6 +3637,7 @@ steps:
# - label: Quantized MoE Test (B200-MI355) # TBD
# timeout_in_minutes: 180
# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
# dind: false
# agent_pool: mi355_1
# working_dir: "/vllm-workspace/"
# source_file_dependencies:
@@ -3646,10 +3663,12 @@ steps:
#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------#
- label: V1 attention (B200-MI355) # TBD
- label: V1 attention (B200-MI355) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
parallelism: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/config/attention.py
@@ -3660,11 +3679,12 @@ steps:
- vllm/envs.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/attention
- pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: V1 Core + KV + Metrics # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3691,6 +3711,7 @@ steps:
- label: V1 Sample + Logits # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
@@ -3711,6 +3732,7 @@ steps:
- label: V1 Spec Decode # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -3724,6 +3746,7 @@ steps:
- label: Weight Loading Multiple GPU # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
@@ -3736,6 +3759,7 @@ steps:
- label: Weight Loading Multiple GPU - Large Models # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
working_dir: "/vllm-workspace/tests"
num_gpus: 2
@@ -3751,6 +3775,7 @@ steps:
- label: Regression # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
.buildkite/test_areas/cuda.yaml
+1
View File
@@ -16,6 +16,7 @@ steps:
commands:
- pytest -v -s cuda/test_cuda_context.py
- pytest -v -s cuda/test_platform_no_cuda_init.py
- pytest -v -s cuda/test_cuda_compatibility_path.py
- label: Cudagraph
device: h200_35gb
.buildkite/test_areas/disaggregated.yaml
+16
View File
@@ -131,6 +131,22 @@ steps:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: NixlConnector PD edge case test (2 GPUs)
key: nixlconnector-pd-edge-cases-2-gpus
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- tests/v1/kv_connector/nixl_integration/
env:
PREFILL_GPU_ID: "0"
DECODE_GPU_ID: "1"
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
timeout_in_minutes: 25
.buildkite/test_areas/engine.yaml
+2
View File
@@ -40,9 +40,11 @@ steps:
source_file_dependencies:
- vllm/v1/engine/
- tests/v1/engine/
- tests/v1/test_tensor_ipc_queue.py
commands:
- pytest -v -s v1/engine/test_preprocess_error_handling.py
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
- pytest -v -s v1/test_tensor_ipc_queue.py
mirror:
amd:
device: mi250_1
.buildkite/test_areas/expert_parallelism.yaml
+1
View File
@@ -52,4 +52,5 @@ steps:
- vllm/compilation/
- tests/distributed/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- pytest -v -s distributed/test_elastic_ep.py
.buildkite/test_areas/fault_tolerance.yaml
+26
View File
@@ -0,0 +1,26 @@
group: Fault Tolerance
depends_on:
- image-build
steps:
- label: Fault Tolerance E2E (2xH100)
key: fault-tolerance-e2e-2xh100
timeout_in_minutes: 35
device: h100
num_devices: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/fault_tolerance/
- vllm/v1/worker/sentinel/
- vllm/entrypoints/serve/fault_tolerance/
- vllm/distributed/elastic_ep/
- vllm/distributed/device_communicators/
- vllm/v1/engine/
- vllm/v1/worker/
- tests/v1/fault_tolerance/
- tests/v1/distributed/test_external_lb_dp.py
commands:
# Base image has no nixl; install it or has_nixl_ep() skips the tests.
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
- pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
.buildkite/test_areas/kernels.yaml
+36
View File
@@ -61,9 +61,45 @@ steps:
source_file_dependencies:
- csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
- vllm/models/deepseek_v4/common/ops/
- vllm/models/deepseek_v4/nvidia/
- tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
- tests/models/test_deepseek_v4_mega_moe.py
commands:
- pytest -v -s kernels/test_fused_deepseek_v4_*.py
- pytest -v -s models/test_deepseek_v4_mega_moe.py
# Catch-all for test files at the tests/kernels root. This job collects
# the whole root so new files are wired by default.
# Files with dedicated jobs elsewhere in this file are excluded via --ignore
# (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in
# their own jobs / Kernels (B200)).
- label: Kernels Root Misc Test (B200)
key: kernels-root-misc-test-b200
timeout_in_minutes: 45
device: b200-k8s
source_file_dependencies:
- csrc/
- vllm/
- tests/kernels/
commands:
- pytest -v -s kernels/
--ignore=kernels/attention
--ignore=kernels/core
--ignore=kernels/helion
--ignore=kernels/ir
--ignore=kernels/mamba
--ignore=kernels/moe
--ignore=kernels/quantization
--ignore=kernels/test_concat_mla_q.py
--ignore=kernels/test_fused_qk_norm_rope_gate.py
--ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
--ignore=kernels/test_top_k_per_row.py
--ignore=kernels/test_kda.py
--ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
--ignore=kernels/test_ll_bf16_gemm.py
--ignore=kernels/test_shuffle_rows.py
# BROKEN on main, pending kernel fixes (B200):
# test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)
- label: Kernels Attention Test %N
key: kernels-attention-test
.buildkite/test_areas/lm_eval.yaml
+2 -2
View File
@@ -337,7 +337,7 @@ steps:
- label: LM Eval KV-Offload (2xH100)
key: kv-offload-medium
timeout_in_minutes: 30
timeout_in_minutes: 45
device: h100
num_devices: 2
source_file_dependencies:
@@ -347,7 +347,7 @@ steps:
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
- label: LM Eval KV-Offload (4xH100)
key: kv-offload-large
.buildkite/test_areas/lora.yaml
+1
View File
@@ -16,6 +16,7 @@ steps:
amd:
dind: false
device: mi300_1
soft_fail: true
working_dir: "/vllm-workspace/tests"
timeout_in_minutes: 85
source_file_dependencies:
.buildkite/test_areas/misc.yaml
+17 -10
View File
@@ -148,6 +148,7 @@ steps:
- pytest -v -s -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s v1/test_kv_cache_spec_registry.py
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
@@ -212,7 +213,7 @@ steps:
- vllm/multimodal
- examples/
commands:
- pip install tensorizer # for tensorizer test
- pip install --no-deps tensorizer # for tensorizer test
# for basic
- python3 basic/offline_inference/chat.py
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -265,6 +266,7 @@ steps:
- vllm/utils/
- vllm/v1/
- tests/v1/tracing
- tests/tracing/
commands:
- "pip install \
'opentelemetry-sdk>=1.26.0' \
@@ -272,6 +274,7 @@ steps:
'opentelemetry-exporter-otlp>=1.26.0' \
'opentelemetry-semantic-conventions-ai>=0.4.1'"
- pytest -v -s v1/tracing
- pytest -v -s tracing
mirror:
amd:
dind: false
@@ -295,6 +298,7 @@ steps:
amd:
device: mi250_1
timeout_in_minutes: 55
soft_fail: true
depends_on:
- image-build-amd
source_file_dependencies:
@@ -394,7 +398,7 @@ steps:
- label: Batch Invariance (A100)
key: batch-invariance-a100
timeout_in_minutes: 40
timeout_in_minutes: 60
device: a100
source_file_dependencies:
- vllm/v1/attention
@@ -404,11 +408,11 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
- label: Batch Invariance (H100)
key: batch-invariance-h100
timeout_in_minutes: 40
timeout_in_minutes: 60
device: h100
source_file_dependencies:
- vllm/v1/attention
@@ -419,12 +423,12 @@ steps:
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
- label: Batch Invariance (B200)
key: batch-invariance-b200
timeout_in_minutes: 35
timeout_in_minutes: 45
device: b200-k8s
source_file_dependencies:
- vllm/v1/attention
@@ -435,11 +439,14 @@ steps:
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
- pytest -v -s v1/determinism/test_matmul_batch_invariant.py
- pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
- pytest -v -s v1/determinism/test_online_batch_invariance.py
- label: Acceptance Length Test (Large Models) # optional
device: h200_35gb
key: acceptance-length-test-large-models
.buildkite/test_areas/model_runner_v2.yaml
+1 -1
View File
@@ -41,7 +41,7 @@ steps:
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pip install tensorizer # for tensorizer test
- pip install --no-deps tensorizer # for tensorizer test
- python3 basic/offline_inference/chat.py # for basic
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
.buildkite/test_areas/models_basic.yaml
+14 -1
View File
@@ -61,6 +61,18 @@ steps:
# FA4 kernel tests require SM100; the suite skips them elsewhere.
- pytest -v -s models/inkling
- label: Kimi K3 Unit Tests (B200)
key: kimi-k3-unit-tests-b200
timeout_in_minutes: 40
device: b200-k8s
source_file_dependencies:
- vllm/models/kimi_k3/
- csrc/libtorch_stable/kimi_k3/
- tests/models/kimi_k3/
commands:
# The native NVIDIA AttnRes kernel requires the SM100 family.
- pytest -v -s models/kimi_k3
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
@@ -70,7 +82,8 @@ steps:
- vllm/
- tests/models/test_utils.py
- tests/models/test_vision.py
- tests/models/test_adapters.py
- tests/models/transformers/fusers/
device: cpu-small
commands:
- pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
- pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/
.buildkite/test_areas/models_language.yaml
-2
View File
@@ -63,7 +63,6 @@ steps:
- tests/models/language/generation
commands:
# Install fast path packages for testing against transformers
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
# Shard the hybrid language model tests that are numerically stable on Hopper.
@@ -105,7 +104,6 @@ steps:
- tests/models/language/generation
commands:
# Install fast path packages for testing against transformers
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
.buildkite/test_areas/spec_decode.yaml
+14
View File
@@ -90,8 +90,10 @@ steps:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
- tests/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
- python3 spec_decode/test_custom_proposer.py
mirror:
amd:
dind: false
@@ -188,3 +190,15 @@ steps:
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
- label: Spec Decode Acceptance Rates Nightly
key: spec-decode-acceptance-rates-nightly
timeout_in_minutes: 60
device: h200_35gb
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
.buildkite/test_areas/torch_abi.yaml
+14
View File
@@ -0,0 +1,14 @@
group: Torch ABI
depends_on:
- image-build
steps:
- label: Torch Stable ABI Audit
key: torch-stable-abi-audit
timeout_in_minutes: 5
source_file_dependencies:
- .buildkite/check-torch-abi.py
- csrc/
- cmake/
- setup.py
commands:
- python3 /vllm-workspace/.buildkite/check-torch-abi.py
.github/mergify.yml
+26
View File
@@ -19,6 +19,7 @@ pull_request_rules:
description: Comment on PR when pre-commit check fails
conditions:
- check-failure=pre-commit
- -check-cancelled=pre-commit
- -closed
- -draft
- or:
@@ -232,6 +233,31 @@ pull_request_rules:
add:
- gpt-oss
- name: label-kimi
description: Automatically apply kimi label
conditions:
- label != stale
- or:
- files~=(?i)kimi
- files~=(?i)moonshot
- title~=(?i)(?:kimi|moonshot)
actions:
label:
add:
- kimi
- name: label-k3
description: Automatically apply k3 label (launch triage; retire after ramp-down)
conditions:
- label != stale
- or:
- files~=(?i)kimi[-_]?k3
- title~=(?i)(?:kimi[-\s]?k3|\bk3\b)
actions:
label:
add:
- k3
- name: label-nvidia
description: Automatically apply nvidia label
conditions:
.github/workflows/issue_autolabel.yml
+19
View File
@@ -130,6 +130,25 @@ jobs:
},
],
},
kimi: {
keywords: [
{ term: "Kimi", searchIn: "both" },
{ term: "Moonshot", searchIn: "both" },
],
substrings: [
{ term: "moonshotai/", searchIn: "both" },
{ term: "kimi", searchIn: "title" },
],
},
k3: {
keywords: [
{ term: "Kimi K3", searchIn: "both" },
{ term: "K3", searchIn: "title" },
],
substrings: [
{ term: "moonshotai/kimi-k3", searchIn: "both" },
],
},
quantization: {
keywords: [
{
.gitignore
-3
View File
@@ -173,9 +173,6 @@ venv.bak/
# mkdocs documentation
/site
docs/argparse
docs/examples/*
!docs/examples/README.md
# mypy
.mypy_cache/
.markdownlint.yaml
+3
View File
@@ -3,6 +3,9 @@ MD007:
MD013: false
MD024:
siblings_only: true
MD025:
# Allow front matter title to be different from the first heading in the document.
front_matter_title: ""
MD031:
list_items: false
MD033: false
.pre-commit-config.yaml
-4
View File
@@ -260,10 +260,6 @@ repos:
files: ^docker/(Dockerfile|versions\.json)$
pass_filenames: false
additional_dependencies: [dockerfile-parse]
- id: attention-backend-docs
name: Check attention backend documentation is up to date
entry: python tools/pre_commit/generate_attention_backend_docs.py --check
language: python
- id: check-boolean-context-manager
name: Check for boolean ops in with-statements
entry: python tools/pre_commit/check_boolean_context_manager.py
CMakeLists.txt
+43 -28
View File
@@ -114,6 +114,11 @@ find_package(Torch REQUIRED)
# Supported NVIDIA architectures.
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
# Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
# targeting it directly.
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
# to support the whole generation without specifying all sub-architectures
@@ -214,10 +219,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
# the set of architectures we want to compile for and remove the from the
# CMAKE_CUDA_FLAGS so that they are not applied globally.
#
# `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch
# as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only
# `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's
# component-specific arch list below.
# `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really
# needs PTX, add `+PTX` to that kernel's component-specific arch list below.
#
clear_cuda_arches(CUDA_ARCH_FLAGS)
extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}")
@@ -227,6 +230,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
cuda_archs_loose_intersection(CUDA_ARCHS
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
if(NOT CUDA_ARCHS)
message(FATAL_ERROR
"No supported CUDA architectures; the build would produce a binary "
"with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; "
"supported: ${CUDA_SUPPORTED_ARCHS}. "
"Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).")
endif()
else()
#
# For other GPU targets override the GPU architectures detected by cmake/torch
@@ -420,7 +430,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
"9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
@@ -695,7 +705,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
@@ -727,23 +737,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"(requires SM90+ and CUDA >= 12.0).")
endif()
# BF16 skinny GEMM (M<=32; weight-bandwidth-bound decode shapes).
# Requires SM90+.
cuda_archs_sm90plus(BF16_SKINNY_GEMM_ARCHS "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND BF16_SKINNY_GEMM_ARCHS)
set(BF16_SKINNY_GEMM_SRCS
"csrc/libtorch_stable/bf16_skinny_gemm_entry.cu"
"csrc/libtorch_stable/bf16_skinny_gemm.cu")
set_gencode_flags_for_srcs(
SRCS "${BF16_SKINNY_GEMM_SRCS}"
CUDA_ARCHS "${BF16_SKINNY_GEMM_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${BF16_SKINNY_GEMM_SRCS}")
message(STATUS "Building bf16_skinny_gemm for archs: ${BF16_SKINNY_GEMM_ARCHS}")
else()
message(STATUS "Not building bf16_skinny_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)
@@ -832,7 +825,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
# require CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -916,7 +909,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -941,7 +934,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# moe_data.cu is used by all CUTLASS MoE kernels.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
@@ -998,7 +991,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
# only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -1064,7 +1057,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# Runtime dispatch is gated in
# vllm/v1/attention/backends/mla/cutlass_mla.py.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
@@ -1086,6 +1079,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set(MLA_ARCHS)
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
"10.0f" "${CUDA_ARCHS}")
endif()
if(KIMI_K3_ATTN_RES_ARCHS)
set(KIMI_K3_ATTN_RES_SRC
"csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${KIMI_K3_ATTN_RES_SRC}"
CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
COMPILE_OPTIONS
"$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
message(STATUS
"Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
endif()
# Hadacore kernels
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
if(HADACORE_ARCHS)
@@ -1127,6 +1138,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
if(KIMI_K3_ATTN_RES_ARCHS)
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
endif()
# Needed by CUTLASS kernels
target_compile_definitions(_C_stable_libtorch PRIVATE
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
benchmarks/attention_benchmarks/benchmark.py
+4
View File
@@ -1358,6 +1358,10 @@ def main():
profile_memory=args.profile_memory,
warmup_ms=args.warmup_ms,
prefill_backend=pb,
kv_lora_rank=args.kv_lora_rank,
qk_nope_head_dim=args.qk_nope_head_dim,
qk_rope_head_dim=args.qk_rope_head_dim,
v_head_dim=args.v_head_dim,
)
result = run_benchmark(config)
benchmarks/kernels/benchmark_inkling_qkvr_prep.py
+176
View File
@@ -0,0 +1,176 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import statistics
import torch
from tabulate import tabulate
from vllm.models.inkling.nvidia.ops import qkvr_prep
from vllm.utils.argparse_utils import FlexibleArgumentParser
def make_inputs(tokens: int, tp_size: int, is_local: bool):
torch.manual_seed(0)
num_q_heads = 64 // tp_size
num_kv_heads = (16 if is_local else 8) // tp_size
head_dim = 128
d_rel = 16
rel_extent = 512 if is_local else 1024
page_size = 16
num_blocks = (tokens + page_size - 1) // page_size
q_width = num_q_heads * head_dim
kv_width = num_kv_heads * head_dim
r_width = num_q_heads * d_rel
device = "cuda"
qkvr = torch.randn(
tokens,
q_width + 2 * kv_width + r_width,
device=device,
dtype=torch.bfloat16,
)
k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
v_weight = torch.randn_like(k_weight)
q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
k_norm_weight = torch.randn_like(q_norm_weight)
rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
conv_cache = torch.zeros(
num_blocks,
num_kv_heads,
page_size,
2 * head_dim,
device=device,
dtype=torch.bfloat16,
)
key_cache = torch.empty(
num_blocks,
page_size,
num_kv_heads,
head_dim,
device=device,
dtype=torch.bfloat16,
)
value_cache = torch.empty_like(key_cache)
positions = torch.arange(tokens, device=device, dtype=torch.int64)
block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
slots = torch.arange(tokens, device=device, dtype=torch.int64)
query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
log_scaling = None
if not is_local:
effective_n = (positions + 1).to(torch.float32)
log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0))
return (
qkvr,
k_weight,
v_weight,
q_norm_weight,
k_norm_weight,
rel_proj,
1e-6,
num_q_heads,
num_kv_heads,
head_dim,
d_rel,
conv_cache,
key_cache,
value_cache,
positions,
block_table,
seq_idx,
slots,
query_start,
slots,
0,
head_dim,
page_size,
log_scaling,
)
def capture(implementation, inputs):
outputs = []
def run():
outputs[:] = implementation.fused_qkvr_prep(*inputs)
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
for _ in range(3):
run()
torch.cuda.current_stream().wait_stream(stream)
torch.accelerator.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
run()
torch.accelerator.synchronize()
return graph, outputs
def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float:
for _ in range(warmup):
graph.replay()
torch.accelerator.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(repeats):
graph.replay()
end.record()
end.synchronize()
return start.elapsed_time(end) * 1000 / repeats
def benchmark(inputs, args) -> float:
graph, _ = capture(qkvr_prep, inputs)
return statistics.median(
time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials)
)
@torch.inference_mode()
def main(args):
rows = []
for tp_size in args.tp_sizes:
for tokens in args.tokens:
for is_local in (True, False):
triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args)
rows.append(
[
tp_size,
tokens,
"local" if is_local else "global",
triton_us,
]
)
print("Inkling QKVR prep (CUDA graph, median latency)")
print(
tabulate(
rows,
headers=[
"TP",
"tokens",
"scope",
"Triton (us)",
],
floatfmt=("d", "d", "", ".2f"),
)
)
if __name__ == "__main__":
parser = FlexibleArgumentParser()
parser.add_argument(
"--tokens",
type=int,
nargs="+",
default=[1 << power for power in range(15)],
)
parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8])
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--repeats", type=int, default=200)
parser.add_argument("--trials", type=int, default=5)
main(parser.parse_args())
benchmarks/kernels/cpu/benchmark_cpu_attn.py
+1 -1
View File
@@ -154,7 +154,7 @@ def main(
scale=scale,
causal=True,
alibi_slopes=None,
sliding_window=window_size,
sliding_window=window_size if sliding_window is not None else -1,
block_table=block_tables,
softcap=0,
scheduler_metadata=metadata,
benchmarks/replayssm/e2e_decode_speedup.py
+267
View File
@@ -0,0 +1,267 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
long greedy decode (CUDA graphs on) once with the standard kernel and once with
ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
separate subprocesses so each gets a clean CUDA context.
The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
CUDA-graph capture on the pre-release Blackwell FP4 path); pass
--no-disable-flashinfer-autotune for non-FP4 models.
Examples:
python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
--model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4
"""
import argparse
import json
import os
import subprocess
import sys
import time
DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
def parse_args():
p = argparse.ArgumentParser(
description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
)
p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
p.add_argument("--prompt", default=DEFAULT_PROMPT)
p.add_argument("--batch-size", type=int, default=256)
p.add_argument("--num-steps", type=int, default=1000)
p.add_argument("--warmup-steps", type=int, default=128)
p.add_argument("--repeats", type=int, default=1)
p.add_argument(
"--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
)
p.add_argument(
"--dtype",
default="bfloat16",
choices=["bfloat16", "float16", "float32", "auto"],
)
p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
p.add_argument("--max-model-len", type=int, default=None)
p.add_argument(
"--disable-flashinfer-autotune",
action=argparse.BooleanOptionalAction,
default=True,
help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
"It is unstable under CUDA-graph capture on the "
"pre-release Blackwell FP4 path; pass "
"--no-disable-flashinfer-autotune for non-FP4 models.",
)
p.add_argument(
"--mamba-ssm-cache-dtype",
default="auto",
choices=["auto", "float32", "float16", "bfloat16"],
help="SSM state dtype (both modes). 'auto' = config-driven; "
"'float32' = fp32 state, 'bfloat16' = s16 state.",
)
p.add_argument(
"--baseline-ssm-config",
default="",
help="Pin the STANDARD baseline's SSM launch config as "
"'bsm,nw' via override_ssm_config (forces the in-process "
"engine so the override reaches the kernel). Empty = off.",
)
p.add_argument(
"--worker",
choices=["standard", "replayssm"],
default=None,
help=argparse.SUPPRESS,
)
return p.parse_args()
def resolve_max_model_len(args) -> int:
if args.max_model_len is not None:
return args.max_model_len
return args.num_steps + 256
def run_worker(args):
# override_ssm_config is a module global; it only reaches the model if the
# engine runs in-process (default V1 spawns a separate EngineCore). Force it.
if args.worker == "standard" and args.baseline_ssm_config:
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
import torch
from vllm import LLM, SamplingParams
mode = args.worker
max_model_len = resolve_max_model_len(args)
llm_kwargs = dict(
model=args.model_id,
tensor_parallel_size=1,
dtype=args.dtype,
max_model_len=max_model_len,
trust_remote_code=True,
enable_prefix_caching=False,
enable_chunked_prefill=False,
max_num_seqs=args.batch_size,
max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
enforce_eager=False,
disable_log_stats=True,
gpu_memory_utilization=args.gpu_memory_utilization,
# SSM state dtype (applies to both standard and ReplaySSM).
mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
)
if args.disable_flashinfer_autotune:
# FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
# re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
if mode == "replayssm":
llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
_ssm_cm = None
if mode == "standard" and args.baseline_ssm_config:
from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
_bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
_ssm_cm = override_ssm_config((_bsm, _nw))
_ssm_cm.__enter__() # active through LLM() graph capture + decode
print(
f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
flush=True,
)
llm = LLM(**llm_kwargs)
prompts = [args.prompt] * args.batch_size
def timed_generate(n_tokens):
sp = SamplingParams(
n=1,
temperature=0.0,
ignore_eos=True,
min_tokens=n_tokens,
max_tokens=n_tokens,
)
if torch.accelerator.is_available():
torch.accelerator.synchronize()
t0 = time.perf_counter()
outs = llm.generate(prompts, sp, use_tqdm=False)
if torch.accelerator.is_available():
torch.accelerator.synchronize()
elapsed = time.perf_counter() - t0
produced = min(len(o.outputs[0].token_ids) for o in outs)
assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
return elapsed
timed_generate(args.warmup_steps)
best = None
for _ in range(args.repeats):
elapsed = timed_generate(args.num_steps)
tok_s = args.batch_size * args.num_steps / elapsed
per_step_ms = elapsed / args.num_steps * 1e3
print(
f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step",
flush=True,
)
if best is None or elapsed < best["elapsed_s"]:
best = {
"mode": mode,
"elapsed_s": elapsed,
"tok_s": tok_s,
"per_step_ms": per_step_ms,
}
print("RESULT_JSON " + json.dumps(best), flush=True)
if _ssm_cm is not None:
_ssm_cm.__exit__(None, None, None)
def run_one_mode(args, mode) -> dict:
cmd = [
sys.executable,
__file__,
"--worker",
mode,
"--model-id",
args.model_id,
"--prompt",
args.prompt,
"--batch-size",
str(args.batch_size),
"--num-steps",
str(args.num_steps),
"--warmup-steps",
str(args.warmup_steps),
"--repeats",
str(args.repeats),
"--buffer-len",
str(args.buffer_len),
"--dtype",
args.dtype,
"--gpu-memory-utilization",
str(args.gpu_memory_utilization),
"--mamba-ssm-cache-dtype",
args.mamba_ssm_cache_dtype,
"--baseline-ssm-config",
args.baseline_ssm_config,
]
cmd.append(
"--disable-flashinfer-autotune"
if args.disable_flashinfer_autotune
else "--no-disable-flashinfer-autotune"
)
if args.max_model_len is not None:
cmd += ["--max-model-len", str(args.max_model_len)]
result = None
proc = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
)
for line in proc.stdout:
sys.stdout.write(line)
sys.stdout.flush()
if line.startswith("RESULT_JSON "):
result = json.loads(line[len("RESULT_JSON ") :])
proc.wait()
if proc.returncode != 0:
raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
if result is None:
raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
return result
def main():
args = parse_args()
if args.worker is not None:
run_worker(args)
return
print(
f"model={args.model_id} batch_size={args.batch_size} "
f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}"
)
std = run_one_mode(args, "standard")
fla = run_one_mode(args, "replayssm")
speedup = std["per_step_ms"] / fla["per_step_ms"]
print()
header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
print(header)
print("-" * len(header))
for r in (std, fla):
print(
f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
)
print("-" * len(header))
print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
if __name__ == "__main__":
main()
cmake/cpu_extension.cmake
+19 -1
View File
@@ -15,6 +15,7 @@ endif()
#
set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
set(ENABLE_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM})
set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
include_directories("${CMAKE_SOURCE_DIR}/csrc")
@@ -96,12 +97,14 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
set(ENABLE_NUMA OFF)
check_sysctl(hw.optional.neon ASIMD_FOUND)
check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND)
else()
find_isa(${CPUINFO} "Power11" POWER11_FOUND)
find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support
find_isa(${CPUINFO} "S390" S390_FOUND)
find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
@@ -111,6 +114,11 @@ else()
set(ARM_BF16_FOUND ON)
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
endif()
if (ENABLE_ARM_I8MM)
set(ARM_I8MM_FOUND ON)
message(STATUS
"ARM I8MM support enabled via VLLM_CPU_ARM_I8MM environment variable")
endif()
# Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
# in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
# overrides the detection result.
@@ -166,6 +174,11 @@ elseif (ASIMD_FOUND)
message(WARNING "BF16 functionality is not available")
set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")
endif()
if(ARM_I8MM_FOUND)
message(STATUS "I8MM extension detected")
string(APPEND MARCH_FLAGS "+i8mm")
add_compile_definitions(ARM_I8MM_SUPPORT)
endif()
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
elseif (S390_FOUND)
message(STATUS "S390 detected")
@@ -447,8 +460,13 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
"csrc/cpu/shm.cpp"
"csrc/cpu/activation_lut_bf16.cpp"
"csrc/cpu/cpu_tanhf_neon.hpp"
"csrc/cpu/cpu_fused_moe.cpp"
${VLLM_EXT_SRC})
if (ARM_BF16_FOUND)
set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC})
if (ARM_I8MM_FOUND)
set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
endif()
endif()
endif()
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
cmake/external_projects/deepgemm.cmake
+3
View File
@@ -68,6 +68,9 @@ endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
endif()
else()
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
endif()
cmake/external_projects/flashmla.cmake
+3 -1
View File
@@ -60,6 +60,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
# CUDA 12.9 has introduced "Family-Specific Architecture Features"
# this supports all compute_10x family
list(APPEND SUPPORT_ARCHS "10.0f")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
list(APPEND SUPPORT_ARCHS "10.7f")
endif()
elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
list(APPEND SUPPORT_ARCHS "10.0a")
endif()
@@ -188,4 +191,3 @@ else()
add_custom_target(_flashmla_C)
add_custom_target(_flashmla_extension_C)
endif()
cmake/external_projects/qutlass.cmake
+5 -1
View File
@@ -55,7 +55,11 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
endif()
else()
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
cmake/utils.cmake
+21 -10
View File
@@ -241,14 +241,15 @@ endmacro()
# `<major>.<minor>`, dedupes them and then sorts them in ascending order and
# stores them in `OUT_ARCHES`.
#
# Example:
# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a"
# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS)
# OUT_ARCHES="7.5;...;9.0"
# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags.
# This handles mismatches such as `arch=compute_20,code=sm_121`.
function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS)
set(_CUDA_ARCHES)
foreach(_ARCH ${CUDA_ARCH_FLAGS})
string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
if (NOT _COMPUTE)
string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
endif()
if (_COMPUTE)
set(_COMPUTE ${CMAKE_MATCH_1})
endif()
@@ -396,14 +397,24 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
# match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
# family. The output uses TGT's value to preserve the user's compilation flags.
set(_CUDA_ARCHS)
# Resolve exact base matches before family fallbacks so a generic entry such
# as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
foreach(_arch ${_SRC_CUDA_ARCHS})
if(_arch MATCHES "[af]$")
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
list(APPEND _CUDA_ARCHS "${_arch}")
endif()
endif()
endforeach()
foreach(_arch ${_SRC_CUDA_ARCHS})
if(_arch MATCHES "[af]$")
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
list(APPEND _CUDA_ARCHS "${_arch}")
elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
list(APPEND _CUDA_ARCHS "${_base}a")
elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
@@ -487,7 +498,7 @@ endfunction()
function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
endif()
csrc/cpu/cpu_arch_macros.h
+11
View File
@@ -172,4 +172,15 @@
#endif // __riscv_v
// Power VSX
#ifdef __powerpc__
// FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which
// implements a vectorised 5-term minimax polynomial using VSX intrinsics.
#define DEFINE_FAST_EXP \
auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
__attribute__((always_inline)) { return vec.exp(); }; \
auto fast_exp_f16 = fast_exp;
#endif // __powerpc__
#endif
csrc/cpu/cpu_fused_moe.cpp
+5 -187
View File
@@ -1,5 +1,6 @@
#include "cpu/cpu_types.hpp"
#include "cpu/utils.hpp"
#include "cpu/cpu_fused_moe_activations.hpp"
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
#include "cpu/cpu_arch_macros.h"
@@ -43,193 +44,9 @@
}()
namespace {
enum class FusedMOEAct {
SiluAndMul,
SwigluOAIAndMul,
GeluAndMul,
GeluTanhAndMul,
};
FusedMOEAct get_act_type(const std::string& act) {
if (act == "silu") {
return FusedMOEAct::SiluAndMul;
} else if (act == "swigluoai") {
return FusedMOEAct::SwigluOAIAndMul;
} else if (act == "gelu") {
return FusedMOEAct::GeluAndMul;
} else if (act == "gelu_tanh") {
return FusedMOEAct::GeluTanhAndMul;
} else {
TORCH_CHECK(false, "Invalid act type: " + act);
}
}
template <typename scalar_t>
void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride,
const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
#if !defined(__aarch64__)
// For GPT-OSS interleaved gate-up weights
alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
16, 18, 20, 22, 24, 26, 28, 30};
vec_op::INT32Vec16 index_vec(index);
#endif
vec_op::FP32Vec16 gate_up_max_vec(7.0);
vec_op::FP32Vec16 up_min_vec(-7.0);
vec_op::FP32Vec16 alpha_vec(1.702);
vec_op::FP32Vec16 one_vec(1.0);
DEFINE_FAST_EXP
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < n_size; n += 32) {
// Note: AdvSIMD does not support gather loads
#if defined(__aarch64__)
vec_op::FP32Vec16 gate_vec(vec_op::uninit);
vec_op::FP32Vec16 up_vec(vec_op::uninit);
vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
#else
vec_op::FP32Vec16 gate_vec(input + n, index_vec);
vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
#endif
gate_vec = gate_vec.min(gate_up_max_vec);
up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
auto glu = gate_vec * sigmoid_vec;
auto gated_output_fp32 = (one_vec + up_vec) * glu;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n / 2);
}
input += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride, const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
DEFINE_FAST_EXP
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
auto silu = gate_vec * sigmoid_vec;
auto gated_output_fp32 = up_vec * silu;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride, const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
vec_op::FP32Vec16 w2_vec(0.5);
alignas(64) float temp[16];
DEFINE_FAST_EXP
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto er_input_vec = gate_vec * w1_vec;
er_input_vec.save(temp);
for (int32_t i = 0; i < 16; ++i) {
temp[i] = std::erf(temp[i]);
}
vec_op::FP32Vec16 er_vec(temp);
auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
auto gated_output_fp32 = up_vec * gelu;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride,
const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
vec_op::FP32Vec16 w2_vec(0.5);
vec_op::FP32Vec16 w3_vec(0.044715);
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
// Note: can't use fast_exp form because diffusiongemma will generate
// wrong results
auto tanh_vec = inner_vec.tanh();
auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
auto gated_output_fp32 = up_vec * gelu_tanh;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
float* __restrict__ input,
scalar_t* __restrict__ output,
const int32_t m, const int32_t n,
const int32_t input_stride,
const int32_t output_stride) {
switch (act) {
case FusedMOEAct::SwigluOAIAndMul:
swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::SiluAndMul:
silu_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::GeluAndMul:
gelu_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::GeluTanhAndMul:
gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
return;
default:
TORCH_CHECK(false, "Unsupported act type.");
}
}
using cpu_fused_moe_utils::apply_gated_act;
using cpu_fused_moe_utils::FusedMOEAct;
template <typename scalar_t, typename gemm_t>
void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
@@ -817,6 +634,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
}
}
}
} // namespace
void prepack_moe_weight(
@@ -864,7 +682,7 @@ void cpu_fused_moe(
const int32_t input_size_2 = w2.size(2);
const int32_t output_size_2 = w2.size(1);
const int32_t topk_num = topk_id.size(1);
const FusedMOEAct act_type = get_act_type(act);
const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
TORCH_CHECK(!skip_weighted || topk_num == 1,
"skip_weighted is only supported for topk=1 on CPU");
csrc/cpu/cpu_fused_moe_activations.hpp
+204
View File
@@ -0,0 +1,204 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP
#define CPU_FUSED_MOE_ACTIVATIONS_HPP
#include <cmath>
#include <cstdint>
#include <string>
#include "cpu/cpu_arch_macros.h"
#include "cpu/utils.hpp"
namespace cpu_fused_moe_utils {
enum class FusedMOEAct {
SiluAndMul,
SwigluOAIAndMul,
GeluAndMul,
GeluTanhAndMul,
};
inline FusedMOEAct get_act_type(const std::string& act) {
if (act == "silu") {
return FusedMOEAct::SiluAndMul;
} else if (act == "swigluoai") {
return FusedMOEAct::SwigluOAIAndMul;
} else if (act == "gelu") {
return FusedMOEAct::GeluAndMul;
} else if (act == "gelu_tanh") {
return FusedMOEAct::GeluTanhAndMul;
} else {
TORCH_CHECK(false, "Invalid act type: " + act);
}
}
template <typename scalar_t>
void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride,
const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
#if !defined(__aarch64__)
// For GPT-OSS interleaved gate-up weights
alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
16, 18, 20, 22, 24, 26, 28, 30};
vec_op::INT32Vec16 index_vec(index);
#endif
vec_op::FP32Vec16 gate_up_max_vec(7.0);
vec_op::FP32Vec16 up_min_vec(-7.0);
vec_op::FP32Vec16 alpha_vec(1.702);
vec_op::FP32Vec16 one_vec(1.0);
DEFINE_FAST_EXP
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < n_size; n += 32) {
// Note: AdvSIMD does not support gather loads
#if defined(__aarch64__)
vec_op::FP32Vec16 gate_vec(vec_op::uninit);
vec_op::FP32Vec16 up_vec(vec_op::uninit);
vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
#else
vec_op::FP32Vec16 gate_vec(input + n, index_vec);
vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
#endif
gate_vec = gate_vec.min(gate_up_max_vec);
up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
auto glu = gate_vec * sigmoid_vec;
auto gated_output_fp32 = (one_vec + up_vec) * glu;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n / 2);
}
input += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride, const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
DEFINE_FAST_EXP
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
auto silu = gate_vec * sigmoid_vec;
auto gated_output_fp32 = up_vec * silu;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride, const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
vec_op::FP32Vec16 w2_vec(0.5);
alignas(64) float temp[16];
DEFINE_FAST_EXP
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto er_input_vec = gate_vec * w1_vec;
er_input_vec.save(temp);
for (int32_t i = 0; i < 16; ++i) {
temp[i] = std::erf(temp[i]);
}
vec_op::FP32Vec16 er_vec(temp);
auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
auto gated_output_fp32 = up_vec * gelu;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
const int32_t m_size, const int32_t n_size,
const int32_t input_stride,
const int32_t output_stride) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
const int32_t dim = n_size / 2;
float* __restrict__ gate = input;
float* __restrict__ up = input + dim;
vec_op::FP32Vec16 one_vec(1.0);
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
vec_op::FP32Vec16 w2_vec(0.5);
vec_op::FP32Vec16 w3_vec(0.044715);
for (int32_t m = 0; m < m_size; ++m) {
for (int32_t n = 0; n < dim; n += 16) {
vec_op::FP32Vec16 gate_vec(gate + n);
vec_op::FP32Vec16 up_vec(up + n);
auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
// Note: can't use fast_exp form because diffusiongemma will generate
// wrong results
auto tanh_vec = inner_vec.tanh();
auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
auto gated_output_fp32 = up_vec * gelu_tanh;
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
gated_output.save(output + n);
}
gate += input_stride;
up += input_stride;
output += output_stride;
}
}
template <typename scalar_t>
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
float* __restrict__ input,
scalar_t* __restrict__ output,
const int32_t m, const int32_t n,
const int32_t input_stride,
const int32_t output_stride) {
switch (act) {
case FusedMOEAct::SwigluOAIAndMul:
swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::SiluAndMul:
silu_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::GeluAndMul:
gelu_and_mul(input, output, m, n, input_stride, output_stride);
return;
case FusedMOEAct::GeluTanhAndMul:
gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
return;
default:
TORCH_CHECK(false, "Unsupported act type.");
}
}
} // namespace cpu_fused_moe_utils
#endif
csrc/cpu/cpu_fused_moe_int8.cpp
+647
View File
@@ -0,0 +1,647 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#include "cpu/cpu_arch_macros.h"
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <optional>
#include <string>
#include "cpu/cpu_fused_moe_activations.hpp"
#include "cpu/cpu_types.hpp"
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
#include "cpu/utils.hpp"
#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT)
#include "cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp"
#define NEON_DISPATCH(SCALAR_TYPE, ...) \
case cpu_utils::ISA::NEON: { \
using gemm_t = \
cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::NEON, SCALAR_TYPE>; \
return __VA_ARGS__(); \
}
#else
#define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON:
#endif
#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \
[&] { \
switch (ISA_TYPE) { \
NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \
default: { \
TORCH_CHECK(false, "Invalid CPU ISA type."); \
} \
} \
}()
namespace {
using cpu_fused_moe_utils::apply_gated_act;
using cpu_fused_moe_utils::FusedMOEAct;
template <typename gemm_t>
void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr,
int8_t* __restrict__ packed_weight_ptr,
const int32_t expert_num,
const int32_t output_size,
const int32_t input_size,
const int64_t expert_stride) {
#pragma omp parallel for
for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) {
gemm_t::pack_weight(weight_ptr + expert_stride * e_idx,
packed_weight_ptr + expert_stride * e_idx, output_size,
input_size);
}
}
// INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp
template <typename scalar_t, typename gemm_t>
void fused_moe_int8_impl(
scalar_t* __restrict__ output, const scalar_t* __restrict__ input,
const int8_t* __restrict__ w13, const int8_t* __restrict__ w2,
const float* __restrict__ w13_scales, const float* __restrict__ w2_scales,
scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias,
const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id,
const FusedMOEAct act_type, const int32_t token_num,
const int32_t expert_num, const int32_t topk_num,
const int32_t input_size_13, const int32_t output_size_13,
const int32_t input_size_2, const int32_t output_size_2,
const bool skip_weighted) {
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0);
TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 0);
TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0);
TORCH_CHECK_EQ(output_size_13 / 2, input_size_2);
const int32_t thread_num = cpu_utils::get_max_threads();
const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
gemm_m_tile_size * input_size_13 * sizeof(int8_t));
const int32_t w2_input_buffer_size =
cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t));
const int32_t w13_n_tile_size = [&]() {
const int64_t cache_size = cpu_utils::get_available_l2_size();
const int32_t n_size_cache_limit =
(cache_size - w13_input_buffer_size) /
(gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t));
const int32_t n_size_thread_limit =
output_size_13 / std::max(1, thread_num / topk_num);
const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>(
std::min(n_size_cache_limit, n_size_thread_limit));
return std::max(n_size, min_w13_n_tile_size);
}();
const int32_t w2_n_tile_size = [&]() {
const int64_t cache_size = cpu_utils::get_available_l2_size();
const int32_t n_size_cache_limit =
(cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_t));
const int32_t n_size_thread_limit =
output_size_2 / std::max(1, thread_num / topk_num);
const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
std::min(n_size_cache_limit, n_size_thread_limit));
return std::max(n_size, gemm_n_tile_size);
}();
int32_t common_buffer_offset = 0;
const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
common_buffer_offset +=
cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
const int32_t expanded_token_num = token_num * topk_num;
const int32_t expand_token_id_buffer_offset = common_buffer_offset;
common_buffer_offset +=
cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
common_buffer_offset +=
cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
const int32_t input_quant_buffer_offset = common_buffer_offset;
common_buffer_offset +=
cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t));
const int32_t input_scale_buffer_offset = common_buffer_offset;
common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float));
const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
common_buffer_offset += cpu_utils::round_up<64>(
expanded_token_num * input_size_2 * sizeof(scalar_t));
const int32_t w13_output_scale_buffer_offset = common_buffer_offset;
common_buffer_offset +=
cpu_utils::round_up<64>(expanded_token_num * sizeof(float));
const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
common_buffer_offset += cpu_utils::round_up<64>(
expanded_token_num * output_size_2 * sizeof(float));
int32_t gemm_thread_buffer_offset = 0;
const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset;
gemm_thread_buffer_offset +=
std::max(w13_input_buffer_size, w2_input_buffer_size);
const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset;
gemm_thread_buffer_offset += cpu_utils::round_up<64>(
gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) *
sizeof(int32_t));
const int32_t ws_output_buffer_offset = 0;
const int32_t ws_thread_buffer_size =
cpu_utils::round_up<64>(output_size_2 * sizeof(float));
const int32_t thread_buffer_size =
std::max(gemm_thread_buffer_offset, ws_thread_buffer_size);
const int32_t buffer_size =
common_buffer_offset + thread_buffer_size * thread_num;
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
uint8_t* common_buffer_start =
cpu_utils::ScratchPadManager::get_scratchpad_manager()
->get_data<uint8_t>();
uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset;
int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>(
common_buffer_start + token_num_per_group_buffer_offset);
int32_t* __restrict__ cu_token_num_per_group_buffer =
reinterpret_cast<int32_t*>(common_buffer_start +
cu_token_num_per_group_buffer_offset);
int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>(
common_buffer_start + expand_token_id_buffer_offset);
int32_t* __restrict__ expand_token_id_index_buffer =
reinterpret_cast<int32_t*>(common_buffer_start +
expand_token_id_index_buffer_offset);
int8_t* __restrict__ input_quant_buffer = reinterpret_cast<int8_t*>(
common_buffer_start + input_quant_buffer_offset);
float* __restrict__ input_scale_buffer =
reinterpret_cast<float*>(common_buffer_start + input_scale_buffer_offset);
std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t));
for (int32_t i = 0; i < expanded_token_num; ++i) {
++token_num_per_group_buffer[topk_id[i]];
}
int32_t token_num_sum = 0;
cu_token_num_per_group_buffer[0] = 0;
int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1;
for (int32_t i = 0; i < expert_num; ++i) {
token_index_buffer[i] = token_num_sum;
token_num_sum += token_num_per_group_buffer[i];
}
for (int32_t i = 0; i < token_num; ++i) {
const int32_t* curr_topk_id = topk_id + i * topk_num;
int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num;
for (int32_t j = 0; j < topk_num; ++j) {
const int32_t curr_expert_id = curr_topk_id[j];
const int32_t curr_index = token_index_buffer[curr_expert_id]++;
expand_token_id_buffer[curr_index] = i;
curr_index_buffer[j] = curr_index;
}
}
// quantize inputs
#pragma omp parallel for
for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
gemm_t::quantize_row(input + token_idx * input_size_13,
input_quant_buffer + token_idx * input_size_13,
input_scale_buffer[token_idx], input_size_13);
}
{
alignas(64) cpu_utils::Counter counter;
cpu_utils::Counter* counter_ptr = &counter;
// w13 GEMM + act
#pragma omp parallel for schedule(static, 1)
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
const int32_t task_num_per_expert =
(output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size;
const int32_t task_num = task_num_per_expert * expert_num;
uint8_t* __restrict__ thread_buffer =
thread_buffer_start + thread_id * thread_buffer_size;
int8_t* __restrict__ gemm_input_buffer =
reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
float* __restrict__ gemm_output_buffer =
reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
common_buffer_start + w13_gemm_output_buffer_offset);
gemm_t gemm;
const int32_t w13_n_group_stride =
gemm_t::WeightOCGroupSize * input_size_13;
const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
for (;;) {
const int32_t task_id = counter_ptr->acquire_counter();
if (task_id >= task_num) {
break;
}
const int32_t curr_expert_id = task_id / task_num_per_expert;
const int32_t curr_output_group_id = task_id % task_num_per_expert;
const int32_t curr_token_num =
token_num_per_group_buffer[curr_expert_id];
if (curr_token_num == 0) {
continue;
}
const int32_t actual_n_tile_size =
std::min(w13_n_tile_size,
output_size_13 - curr_output_group_id * w13_n_tile_size);
const int32_t* __restrict__ curr_expand_token_id_buffer =
expand_token_id_buffer +
cu_token_num_per_group_buffer[curr_expert_id];
scalar_t* __restrict__ curr_w13_gemm_output_buffer =
w13_gemm_output_buffer +
cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 +
curr_output_group_id * w13_n_tile_size / 2;
const int8_t* w13_weight_ptr_0 = nullptr;
const int8_t* w13_weight_ptr_1 = nullptr;
const float* w13_scale_ptr_0 = nullptr;
const float* w13_scale_ptr_1 = nullptr;
scalar_t* w13_bias_ptr_0 = nullptr;
scalar_t* w13_bias_ptr_1 = nullptr;
if (act_type == FusedMOEAct::SwigluOAIAndMul) {
const int32_t output_offset = curr_output_group_id * w13_n_tile_size;
w13_weight_ptr_0 = w13 +
curr_expert_id * input_size_13 * output_size_13 +
output_offset * input_size_13;
w13_weight_ptr_1 =
w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13;
w13_scale_ptr_0 =
w13_scales + curr_expert_id * output_size_13 + output_offset;
w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2;
if (w13_bias != nullptr) {
w13_bias_ptr_0 =
w13_bias + curr_expert_id * output_size_13 + output_offset;
w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2;
}
} else {
const int32_t output_offset =
curr_output_group_id * (w13_n_tile_size / 2);
w13_weight_ptr_0 = w13 +
curr_expert_id * input_size_13 * output_size_13 +
output_offset * input_size_13;
w13_weight_ptr_1 =
w13_weight_ptr_0 + output_size_13 / 2 * input_size_13;
w13_scale_ptr_0 =
w13_scales + curr_expert_id * output_size_13 + output_offset;
w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2;
if (w13_bias != nullptr) {
w13_bias_ptr_0 =
w13_bias + curr_expert_id * output_size_13 + output_offset;
w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2;
}
}
for (int32_t token_idx = 0; token_idx < curr_token_num;
token_idx += gemm_m_tile_size) {
const int32_t actual_token_num =
std::min(gemm_m_tile_size, curr_token_num - token_idx);
const int8_t* input_rows[gemm_m_tile_size];
alignas(64) float input_scales[gemm_m_tile_size];
// gather and pack
for (int32_t i = 0; i < actual_token_num; ++i) {
const int32_t curr_token_id = curr_expand_token_id_buffer[i];
input_rows[i] = input_quant_buffer + curr_token_id * input_size_13;
input_scales[i] = input_scale_buffer[curr_token_id];
}
gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
actual_token_num, input_size_13);
curr_expand_token_id_buffer += actual_token_num;
const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0;
const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1;
const float* w13_scale_ptr_0_iter = w13_scale_ptr_0;
const float* w13_scale_ptr_1_iter = w13_scale_ptr_1;
scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0;
scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1;
float* w13_output_buffer_0_iter = gemm_output_buffer;
float* w13_output_buffer_1_iter =
gemm_output_buffer + actual_n_tile_size / 2;
for (int32_t i = 0; i < actual_n_tile_size;
i += min_w13_n_tile_size) {
auto* output_0_int32 =
reinterpret_cast<int32_t*>(w13_output_buffer_0_iter);
gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32,
actual_token_num, input_size_13, w13_n_group_stride,
actual_n_tile_size);
gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter,
input_scales, w13_scale_ptr_0_iter,
actual_token_num, gemm_n_tile_size,
actual_n_tile_size);
if (w13_bias != nullptr) {
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
w13_output_buffer_0_iter, w13_output_buffer_0_iter,
w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size,
actual_n_tile_size);
w13_bias_ptr_0_iter += gemm_n_tile_size;
}
auto* output_1_int32 =
reinterpret_cast<int32_t*>(w13_output_buffer_1_iter);
gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32,
actual_token_num, input_size_13, w13_n_group_stride,
actual_n_tile_size);
gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter,
input_scales, w13_scale_ptr_1_iter,
actual_token_num, gemm_n_tile_size,
actual_n_tile_size);
if (w13_bias != nullptr) {
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
w13_output_buffer_1_iter, w13_output_buffer_1_iter,
w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size,
actual_n_tile_size);
w13_bias_ptr_1_iter += gemm_n_tile_size;
}
w13_weight_ptr_0_iter += w13_n_tile_stride;
w13_weight_ptr_1_iter += w13_n_tile_stride;
w13_scale_ptr_0_iter += gemm_n_tile_size;
w13_scale_ptr_1_iter += gemm_n_tile_size;
w13_output_buffer_0_iter += gemm_n_tile_size;
w13_output_buffer_1_iter += gemm_n_tile_size;
}
apply_gated_act(act_type, gemm_output_buffer,
curr_w13_gemm_output_buffer, actual_token_num,
actual_n_tile_size, actual_n_tile_size, input_size_2);
curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
}
}
}
}
auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
common_buffer_start + w13_gemm_output_buffer_offset);
float* __restrict__ w13_output_scale_buffer = reinterpret_cast<float*>(
common_buffer_start + w13_output_scale_buffer_offset);
// quantize w2 inputs - in place
#pragma omp parallel for
for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) {
scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2;
int8_t* output_row = reinterpret_cast<int8_t*>(input_row);
gemm_t::quantize_row(input_row, output_row,
w13_output_scale_buffer[token_idx], input_size_2);
}
{
alignas(64) cpu_utils::Counter counter;
cpu_utils::Counter* counter_ptr = &counter;
// w2 gemm
#pragma omp parallel for schedule(static, 1)
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
const int32_t task_num_per_expert =
(output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size;
const int32_t task_num = task_num_per_expert * expert_num;
uint8_t* __restrict__ thread_buffer =
thread_buffer_start + thread_id * thread_buffer_size;
int8_t* __restrict__ gemm_input_buffer =
reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
float* __restrict__ gemm_output_buffer =
reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
common_buffer_start + w2_gemm_output_buffer_offset);
gemm_t gemm;
const int32_t w2_n_group_stride =
gemm_t::WeightOCGroupSize * input_size_2;
const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
for (;;) {
const int32_t task_id = counter_ptr->acquire_counter();
if (task_id >= task_num) {
break;
}
const int32_t curr_expert_id = task_id / task_num_per_expert;
const int32_t curr_output_group_id = task_id % task_num_per_expert;
const int32_t curr_token_num =
token_num_per_group_buffer[curr_expert_id];
if (curr_token_num == 0) {
continue;
}
const int32_t actual_n_tile_size =
std::min(w2_n_tile_size,
output_size_2 - curr_output_group_id * w2_n_tile_size);
scalar_t* __restrict__ curr_w13_gemm_output_buffer =
w13_gemm_output_buffer +
cu_token_num_per_group_buffer[curr_expert_id] * input_size_2;
float* __restrict__ curr_w13_output_scale_buffer =
w13_output_scale_buffer +
cu_token_num_per_group_buffer[curr_expert_id];
float* __restrict__ curr_w2_gemm_output_buffer =
w2_gemm_output_buffer +
cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 +
curr_output_group_id * w2_n_tile_size;
const int8_t* __restrict__ w2_weight_ptr =
w2 + curr_expert_id * output_size_2 * input_size_2 +
curr_output_group_id * w2_n_tile_size * input_size_2;
const float* __restrict__ w2_scale_ptr =
w2_scales + curr_expert_id * output_size_2 +
curr_output_group_id * w2_n_tile_size;
scalar_t* w2_bias_ptr = nullptr;
if (w2_bias != nullptr) {
w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 +
curr_output_group_id * w2_n_tile_size;
}
for (int32_t token_idx = 0; token_idx < curr_token_num;
token_idx += gemm_m_tile_size) {
const int32_t actual_token_num =
std::min(gemm_m_tile_size, curr_token_num - token_idx);
const int8_t* input_rows[gemm_m_tile_size];
alignas(64) float input_scales[gemm_m_tile_size];
for (int32_t i = 0; i < actual_token_num; ++i) {
input_rows[i] = reinterpret_cast<const int8_t*>(
curr_w13_gemm_output_buffer + i * input_size_2);
input_scales[i] = curr_w13_output_scale_buffer[i];
}
gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
actual_token_num, input_size_2);
const int8_t* w2_weight_ptr_iter = w2_weight_ptr;
const float* w2_scale_ptr_iter = w2_scale_ptr;
scalar_t* w2_bias_ptr_iter = w2_bias_ptr;
float* curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer;
for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
auto* output_int32 = reinterpret_cast<int32_t*>(gemm_output_buffer);
gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32,
actual_token_num, input_size_2, w2_n_group_stride,
gemm_n_tile_size);
gemm_t::dequantize_tile(output_int32, gemm_output_buffer,
input_scales, w2_scale_ptr_iter,
actual_token_num, gemm_n_tile_size,
gemm_n_tile_size);
if (w2_bias != nullptr) {
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter,
actual_token_num, gemm_n_tile_size, gemm_n_tile_size);
w2_bias_ptr_iter += gemm_n_tile_size;
}
for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) {
std::memcpy(
curr_w2_gemm_output_buffer_iter + m_idx * output_size_2,
gemm_output_buffer + m_idx * gemm_n_tile_size,
gemm_n_tile_size * sizeof(float));
}
w2_weight_ptr_iter += w2_n_tile_stride;
w2_scale_ptr_iter += gemm_n_tile_size;
curr_w2_gemm_output_buffer_iter += gemm_n_tile_size;
}
curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
curr_w13_output_scale_buffer += gemm_m_tile_size;
curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2;
}
}
}
}
{
alignas(64) cpu_utils::Counter counter;
cpu_utils::Counter* counter_ptr = &counter;
#pragma omp parallel for schedule(static, 1)
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
uint8_t* __restrict__ thread_buffer =
thread_buffer_start + thread_id * thread_buffer_size;
float* __restrict__ ws_output_buffer =
reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset);
float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
common_buffer_start + w2_gemm_output_buffer_offset);
for (;;) {
const int32_t token_id = counter_ptr->acquire_counter();
if (token_id >= token_num) {
break;
}
int32_t* __restrict__ curr_expand_token_id_index_buffer =
expand_token_id_index_buffer + token_id * topk_num;
const float* __restrict__ curr_weight =
topk_weights + token_id * topk_num;
const float first_weight = skip_weighted ? 1.0f : curr_weight[0];
scalar_t* __restrict__ curr_output_buffer =
output + token_id * output_size_2;
if (topk_num > 1) {
int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
float* w2_output_iter =
w2_gemm_output_buffer + w2_output_idx * output_size_2;
float* ws_output_buffer_iter = ws_output_buffer;
vec_op::FP32Vec16 weight_vec(first_weight);
for (int32_t i = 0; i < output_size_2; i += 16) {
vec_op::FP32Vec16 vec(w2_output_iter);
(vec * weight_vec).save(ws_output_buffer_iter);
w2_output_iter += 16;
ws_output_buffer_iter += 16;
}
for (int32_t idx = 1; idx < topk_num - 1; ++idx) {
w2_output_idx = curr_expand_token_id_index_buffer[idx];
w2_output_iter =
w2_gemm_output_buffer + w2_output_idx * output_size_2;
ws_output_buffer_iter = ws_output_buffer;
weight_vec = vec_op::FP32Vec16(curr_weight[idx]);
for (int32_t i = 0; i < output_size_2; i += 16) {
vec_op::FP32Vec16 vec(w2_output_iter);
vec_op::FP32Vec16 sum(ws_output_buffer_iter);
(sum + vec * weight_vec).save(ws_output_buffer_iter);
w2_output_iter += 16;
ws_output_buffer_iter += 16;
}
}
const int32_t last_idx = topk_num - 1;
w2_output_idx = curr_expand_token_id_index_buffer[last_idx];
w2_output_iter =
w2_gemm_output_buffer + w2_output_idx * output_size_2;
ws_output_buffer_iter = ws_output_buffer;
scalar_t* curr_output_buffer_iter = curr_output_buffer;
weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]);
for (int32_t i = 0; i < output_size_2; i += 16) {
vec_op::FP32Vec16 vec(w2_output_iter);
vec_op::FP32Vec16 sum(ws_output_buffer_iter);
scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter);
w2_output_iter += 16;
ws_output_buffer_iter += 16;
curr_output_buffer_iter += 16;
}
} else {
const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
float* w2_output_iter =
w2_gemm_output_buffer + w2_output_idx * output_size_2;
scalar_t* curr_output_buffer_iter = curr_output_buffer;
vec_op::FP32Vec16 weight_vec(first_weight);
for (int32_t i = 0; i < output_size_2; i += 16) {
vec_op::FP32Vec16 vec(w2_output_iter);
scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter);
w2_output_iter += 16;
curr_output_buffer_iter += 16;
}
}
}
}
}
}
} // namespace
void prepack_moe_weight_int8(
const torch::Tensor& weight, // [expert_num, output_size, input_size]
torch::Tensor& packed_weight, const std::string& isa) {
TORCH_CHECK(weight.is_contiguous());
const int32_t expert_num = weight.size(0);
const int32_t output_size = weight.size(1);
const int32_t input_size = weight.size(2);
const int64_t expert_stride = weight.stride(0);
const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
TORCH_CHECK_EQ(output_size % 32, 0);
CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() {
TORCH_CHECK_EQ(input_size % gemm_t::K, 0);
prepack_moe_weight_int8_impl<gemm_t>(
weight.data_ptr<int8_t>(), packed_weight.data_ptr<int8_t>(), expert_num,
output_size, input_size, expert_stride);
});
}
void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
const torch::Tensor& w13, const torch::Tensor& w2,
const torch::Tensor& w13_scale,
const torch::Tensor& w2_scale,
const std::optional<torch::Tensor>& w13_bias,
const std::optional<torch::Tensor>& w2_bias,
const torch::Tensor& topk_weights,
const torch::Tensor& topk_id, const bool skip_weighted,
const std::string& act, const std::string& isa) {
const int32_t token_num = input.size(0);
const int32_t input_size_13 = input.size(1);
const int64_t input_stride = input.stride(0);
TORCH_CHECK_EQ(input_stride, input_size_13);
const int32_t expert_num = w13.size(0);
const int32_t output_size_13 = w13.size(1);
const int32_t input_size_2 = w2.size(2);
const int32_t output_size_2 = w2.size(1);
const int32_t topk_num = topk_id.size(1);
const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
TORCH_CHECK(!skip_weighted || topk_num == 1,
"skip_weighted is only supported for topk=1 on CPU");
VLLM_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "cpu_fused_moe_int8", [&]() {
CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() {
fused_moe_int8_impl<scalar_t, gemm_t>(
output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
w13.data_ptr<int8_t>(), w2.data_ptr<int8_t>(),
w13_scale.data_ptr<float>(), w2_scale.data_ptr<float>(),
w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr,
w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(),
act_type, token_num, expert_num, topk_num, input_size_13,
output_size_13, input_size_2, output_size_2, skip_weighted);
});
});
}
csrc/cpu/cpu_types_vsx.hpp
+12 -3
View File
@@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
explicit FP32Vec4(__vector float data) : reg(data) {}
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
};
struct FP32Vec8 : public Vec<FP32Vec8> {
@@ -316,7 +316,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
explicit FP32Vec8(f32x4x2_t data) : reg(data) {}
explicit FP32Vec8(const FP32Vec8& data) {
FP32Vec8(const FP32Vec8& data) {
reg.val[0] = data.reg.val[0];
reg.val[1] = data.reg.val[1];
}
@@ -593,7 +593,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
explicit FP32Vec16(const FP32Vec16& data) {
FP32Vec16(const FP32Vec16& data) {
reg.val[0] = data.reg.val[0];
reg.val[1] = data.reg.val[1];
reg.val[2] = data.reg.val[2];
@@ -747,6 +747,15 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
}
FP32Vec16 exp() const {
FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
auto lo_e = lo.exp();
auto hi_e = hi.exp();
return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
hi_e.reg.val[0], hi_e.reg.val[1]});
}
float reduce_max() {
__vector float max01 = vec_max(reg.val[0], reg.val[1]);
__vector float max23 = vec_max(reg.val[2], reg.val[3]);
csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp
+3
View File
@@ -31,6 +31,9 @@ class MicroGemm {
}
};
template <cpu_utils::ISA isa, typename scalar_t>
class MicroGemmINT8;
template <int32_t n_size, typename scalar_t>
FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
scalar_t* __restrict__ d_ptr,
csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp
+424
View File
@@ -0,0 +1,424 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#ifndef CPU_MICRO_GEMM_INT8_NEON_HPP
#define CPU_MICRO_GEMM_INT8_NEON_HPP
#include <algorithm>
#include <cstdint>
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
#include <arm_bf16.h>
#include <arm_neon.h>
#include <c10/util/BFloat16.h>
#include <c10/util/Exception.h>
#include <c10/util/Half.h>
namespace cpu_micro_gemm {
namespace neon_smmla {
constexpr int32_t K = 8;
constexpr int32_t Cols = 2;
constexpr int32_t TileSize = K * Cols;
FORCE_INLINE float32x4x2_t load_as_f32(const float* input) {
float32x4x2_t result;
result.val[0] = vld1q_f32(input);
result.val[1] = vld1q_f32(input + 4);
return result;
}
FORCE_INLINE float32x4x2_t load_as_f32(const c10::Half* input) {
const auto input_vec = vld1q_f16(reinterpret_cast<const float16_t*>(input));
float32x4x2_t result;
result.val[0] = vcvt_f32_f16(vget_low_f16(input_vec));
result.val[1] = vcvt_f32_f16(vget_high_f16(input_vec));
return result;
}
FORCE_INLINE float32x4x2_t load_as_f32(const c10::BFloat16* input) {
const auto input_vec = vld1q_bf16(reinterpret_cast<const bfloat16_t*>(input));
float32x4x2_t result;
result.val[0] = vcvt_f32_bf16(vget_low_bf16(input_vec));
result.val[1] = vcvt_f32_bf16(vget_high_bf16(input_vec));
return result;
}
FORCE_INLINE void store_acc_rowpair(const int32x4_t acc01,
const int32x4_t acc23,
const int32x4_t acc45,
const int32x4_t acc67,
int32_t* __restrict__ c_ptr,
const int64_t ldc, const int32_t m_rows) {
if (m_rows == 0) {
return;
}
vst1q_s32(c_ptr, vcombine_s32(vget_low_s32(acc01), vget_low_s32(acc23)));
vst1q_s32(c_ptr + 4, vcombine_s32(vget_low_s32(acc45), vget_low_s32(acc67)));
if (m_rows == 2) {
vst1q_s32(c_ptr + ldc,
vcombine_s32(vget_high_s32(acc01), vget_high_s32(acc23)));
vst1q_s32(c_ptr + ldc + 4,
vcombine_s32(vget_high_s32(acc45), vget_high_s32(acc67)));
}
}
FORCE_INLINE void gemm_micro_smmla_8x8_packed_a(
const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
const int64_t ldc) {
const int32x4_t zero = vdupq_n_s32(0);
int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
int32x4_t acc4501 = zero, acc4523 = zero, acc4545 = zero, acc4567 = zero;
int32x4_t acc6701 = zero, acc6723 = zero, acc6745 = zero, acc6767 = zero;
const int8_t* __restrict__ a_tile = a_packed;
const int8_t* __restrict__ b_tile = b_packed;
#pragma GCC unroll 8
for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
const int8x16_t a_tile01 = vld1q_s8(a_tile);
const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
const int8x16_t a_tile45 = vld1q_s8(a_tile + 2 * TileSize);
const int8x16_t a_tile67 = vld1q_s8(a_tile + 3 * TileSize);
const int8x16_t b_tile01 = vld1q_s8(b_tile);
const int8x16_t b_tile23 = vld1q_s8(b_tile + TileSize);
const int8x16_t b_tile45 = vld1q_s8(b_tile + 2 * TileSize);
const int8x16_t b_tile67 = vld1q_s8(b_tile + 3 * TileSize);
acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
acc4501 = vmmlaq_s32(acc4501, a_tile45, b_tile01);
acc6701 = vmmlaq_s32(acc6701, a_tile67, b_tile01);
acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
acc4523 = vmmlaq_s32(acc4523, a_tile45, b_tile23);
acc6723 = vmmlaq_s32(acc6723, a_tile67, b_tile23);
acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
acc4545 = vmmlaq_s32(acc4545, a_tile45, b_tile45);
acc6745 = vmmlaq_s32(acc6745, a_tile67, b_tile45);
acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
acc4567 = vmmlaq_s32(acc4567, a_tile45, b_tile67);
acc6767 = vmmlaq_s32(acc6767, a_tile67, b_tile67);
a_tile += 4 * TileSize;
b_tile += 4 * TileSize;
}
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
std::min(2, m));
store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
std::min(2, std::max(0, m - 2)));
store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
std::min(2, std::max(0, m - 4)));
store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
std::min(2, std::max(0, m - 6)));
}
FORCE_INLINE void gemm_micro_smmla_4x16_packed_a(
const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
const int64_t b_n_group_stride, const int64_t ldc) {
const int32_t m_rows_01 = std::min(2, m);
const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
const int32x4_t zero = vdupq_n_s32(0);
int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
int32x4_t acc0189 = zero, acc011011 = zero, acc011213 = zero,
acc011415 = zero;
int32x4_t acc2389 = zero, acc231011 = zero, acc231213 = zero,
acc231415 = zero;
const int8_t* __restrict__ a_tile = a_packed;
// note: b packs 8 panels contiguously, so we need 2 b_tile ptrs
// for the 4x16 microkernel
const int8_t* __restrict__ b_tile0 = b_packed;
const int8_t* __restrict__ b_tile1 = b_packed + b_n_group_stride;
#pragma GCC unroll 8
for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
const int8x16_t a_tile01 = vld1q_s8(a_tile);
const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
const int8x16_t b_tile01 = vld1q_s8(b_tile0);
const int8x16_t b_tile23 = vld1q_s8(b_tile0 + TileSize);
const int8x16_t b_tile45 = vld1q_s8(b_tile0 + 2 * TileSize);
const int8x16_t b_tile67 = vld1q_s8(b_tile0 + 3 * TileSize);
const int8x16_t b_tile89 = vld1q_s8(b_tile1);
const int8x16_t b_tile1011 = vld1q_s8(b_tile1 + TileSize);
const int8x16_t b_tile1213 = vld1q_s8(b_tile1 + 2 * TileSize);
const int8x16_t b_tile1415 = vld1q_s8(b_tile1 + 3 * TileSize);
acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
acc0189 = vmmlaq_s32(acc0189, a_tile01, b_tile89);
acc2389 = vmmlaq_s32(acc2389, a_tile23, b_tile89);
acc011011 = vmmlaq_s32(acc011011, a_tile01, b_tile1011);
acc231011 = vmmlaq_s32(acc231011, a_tile23, b_tile1011);
acc011213 = vmmlaq_s32(acc011213, a_tile01, b_tile1213);
acc231213 = vmmlaq_s32(acc231213, a_tile23, b_tile1213);
acc011415 = vmmlaq_s32(acc011415, a_tile01, b_tile1415);
acc231415 = vmmlaq_s32(acc231415, a_tile23, b_tile1415);
a_tile += 2 * TileSize;
b_tile0 += 4 * TileSize;
b_tile1 += 4 * TileSize;
}
// rows 0-1, columns 0-7
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
// rows 0-1, columns 8-15
store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
m_rows_01);
// rows 2-3, columns 0-7
store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
m_rows_23);
// rows 2-3, columns 8-15
store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
c_ptr + 2 * ldc + 8, ldc, m_rows_23);
}
} // namespace neon_smmla
template <typename scalar_t>
class MicroGemmINT8<cpu_utils::ISA::NEON, scalar_t> {
public:
static constexpr int32_t K = neon_smmla::K;
static constexpr int32_t Mr = 8;
static constexpr int32_t Nr = 8;
static constexpr int32_t NrGemv = 16;
static constexpr int32_t MaxMSize = 8;
static constexpr int32_t NSize = 32;
static constexpr int32_t WeightOCGroupSize = Nr;
static_assert(MaxMSize % Mr == 0);
static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output,
float& scale, const int32_t size) {
TORCH_CHECK_EQ(size % K, 0);
float32x4_t max_vec = vdupq_n_f32(0.0f);
for (int32_t i = 0; i < size; i += K) {
const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[0]));
max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[1]));
}
const float abs_max = std::max(vmaxvq_f32(max_vec), 1.0e-7f);
scale = abs_max / 127.0f;
const float32x4_t inv_scale_vec = vdupq_n_f32(127.0f / abs_max);
for (int32_t i = 0; i < size; i += K) {
const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
const int32x4_t output_low =
vcvtnq_s32_f32(vmulq_f32(input_vec.val[0], inv_scale_vec));
const int32x4_t output_high =
vcvtnq_s32_f32(vmulq_f32(input_vec.val[1], inv_scale_vec));
const int16x8_t output_s16 =
vcombine_s16(vqmovn_s32(output_low), vqmovn_s32(output_high));
vst1_s8(output + i, vqmovn_s16(output_s16));
}
}
// with current code, fusing this into the gemm micro kernel didn't move the
// needle
static FORCE_INLINE void dequantize_tile(
int32_t* input, float* output, const float* __restrict__ input_scales,
const float* __restrict__ weight_scales, const int32_t m, const int32_t n,
const int32_t stride) {
TORCH_CHECK_EQ(n % 4, 0);
for (int32_t m_idx = 0; m_idx < m; ++m_idx) {
const float32x4_t input_scale_vec = vdupq_n_f32(input_scales[m_idx]);
for (int32_t n_idx = 0; n_idx < n; n_idx += 4) {
const int32x4_t input_vec = vld1q_s32(input + m_idx * stride + n_idx);
const float32x4_t weight_scale_vec = vld1q_f32(weight_scales + n_idx);
const float32x4_t output_vec =
vmulq_f32(vcvtq_f32_s32(input_vec),
vmulq_f32(input_scale_vec, weight_scale_vec));
vst1q_f32(output + m_idx * stride + n_idx, output_vec);
}
}
}
// physical layout [
// M / (8 or 4); Mr is 8 or 4
// K / 8; K for smmla is 8
// 4, ; 4 row-pairs for each 8 rows
// 2, ; row-pair is 2 rows
// 4 ; 4 elements per row
// ]
static void pack_input_from_rows(const int8_t* const* __restrict__ rows,
int8_t* __restrict__ a_packed,
const int32_t m, const int32_t k) {
TORCH_CHECK(m > 0 && m <= MaxMSize);
TORCH_CHECK(k % K == 0);
const int8x8_t zero = vdup_n_s8(0);
for (int32_t row_base = 0; row_base < m; row_base += Mr) {
const int32_t panel_m = std::min(Mr, m - row_base);
const int8_t* const* panel_rows = rows + row_base;
int8_t* __restrict__ out = a_packed + row_base * k;
// fast path for full 8-row panels (fast path for 4-row panels didn't move
// the needle)
if (panel_m == Mr) {
const int8_t* __restrict__ row0 = panel_rows[0];
const int8_t* __restrict__ row1 = panel_rows[1];
const int8_t* __restrict__ row2 = panel_rows[2];
const int8_t* __restrict__ row3 = panel_rows[3];
const int8_t* __restrict__ row4 = panel_rows[4];
const int8_t* __restrict__ row5 = panel_rows[5];
const int8_t* __restrict__ row6 = panel_rows[6];
const int8_t* __restrict__ row7 = panel_rows[7];
int32_t k_idx = 0;
for (; k_idx + 2 * K <= k; k_idx += 2 * K) {
int8_t* __restrict__ block0 = out;
int8_t* __restrict__ block1 = out + 4 * neon_smmla::TileSize;
int8x16_t a0 = vld1q_s8(row0 + k_idx);
int8x16_t a1 = vld1q_s8(row1 + k_idx);
vst1q_s8(block0, vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
vst1q_s8(block1, vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
a0 = vld1q_s8(row2 + k_idx);
a1 = vld1q_s8(row3 + k_idx);
vst1q_s8(block0 + neon_smmla::TileSize,
vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
vst1q_s8(block1 + neon_smmla::TileSize,
vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
a0 = vld1q_s8(row4 + k_idx);
a1 = vld1q_s8(row5 + k_idx);
vst1q_s8(block0 + 2 * neon_smmla::TileSize,
vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
vst1q_s8(block1 + 2 * neon_smmla::TileSize,
vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
a0 = vld1q_s8(row6 + k_idx);
a1 = vld1q_s8(row7 + k_idx);
vst1q_s8(block0 + 3 * neon_smmla::TileSize,
vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
vst1q_s8(block1 + 3 * neon_smmla::TileSize,
vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
out += 8 * neon_smmla::TileSize;
}
for (; k_idx < k; k_idx += K) {
int8x8_t a0 = vld1_s8(row0 + k_idx);
int8x8_t a1 = vld1_s8(row1 + k_idx);
vst1q_s8(out, vcombine_s8(a0, a1));
a0 = vld1_s8(row2 + k_idx);
a1 = vld1_s8(row3 + k_idx);
vst1q_s8(out + neon_smmla::TileSize, vcombine_s8(a0, a1));
a0 = vld1_s8(row4 + k_idx);
a1 = vld1_s8(row5 + k_idx);
vst1q_s8(out + 2 * neon_smmla::TileSize, vcombine_s8(a0, a1));
a0 = vld1_s8(row6 + k_idx);
a1 = vld1_s8(row7 + k_idx);
vst1q_s8(out + 3 * neon_smmla::TileSize, vcombine_s8(a0, a1));
out += 4 * neon_smmla::TileSize;
}
continue;
}
const int32_t row_pairs = (panel_m <= 4) ? 2 : Mr / 2;
for (int32_t k_idx = 0; k_idx < k; k_idx += K) {
for (int32_t pair_idx = 0; pair_idx < row_pairs; ++pair_idx) {
const int32_t row_idx = pair_idx * 2;
const int8x8_t row0 =
(row_idx < panel_m) ? vld1_s8(panel_rows[row_idx] + k_idx) : zero;
const int8x8_t row1 = (row_idx + 1 < panel_m)
? vld1_s8(panel_rows[row_idx + 1] + k_idx)
: zero;
vst1q_s8(out, vcombine_s8(row0, row1));
out += neon_smmla::TileSize;
}
}
}
}
// physical layout [
// N / 8; Nr is 8
// K / 8; K for smmla is 8
// 4, ; 4 col-pairs for each 8 cols
// 2, ; col-pair is 2 cols
// 4 ; 4 elements per col
// ]
static void pack_weight(const int8_t* __restrict__ weight,
int8_t* __restrict__ packed_weight,
const int32_t output_size, const int32_t input_size) {
TORCH_CHECK(output_size % NSize == 0);
TORCH_CHECK(input_size % K == 0);
for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
int8_t* __restrict__ dst = packed_weight + o_idx * input_size;
for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) {
for (int32_t pair_idx = 0; pair_idx < Nr;
pair_idx += neon_smmla::Cols) {
const int8_t* __restrict__ row0 =
weight + (o_idx + pair_idx) * input_size + k_idx;
const int8_t* __restrict__ row1 = row0 + input_size;
vst1q_s8(dst, vcombine_s8(vld1_s8(row0), vld1_s8(row1)));
dst += neon_smmla::TileSize;
}
}
}
}
void gemm(const int8_t* __restrict__ a_packed,
const int8_t* __restrict__ b_packed, int32_t* __restrict__ c,
const int32_t m, const int32_t k, const int64_t b_n_group_stride,
const int64_t ldc) const {
TORCH_CHECK(m > 0 && m <= MaxMSize);
TORCH_CHECK(k % K == 0);
for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
const int8_t* __restrict__ b_panel = b_packed + n_idx * k;
for (int32_t row_base = 0; row_base < m; row_base += Mr) {
const int32_t panel_m = std::min(Mr, m - row_base);
const int8_t* __restrict__ a_panel = a_packed + row_base * k;
int32_t* __restrict__ c_panel = c + row_base * ldc + n_idx;
if (panel_m <= 4) {
neon_smmla::gemm_micro_smmla_4x16_packed_a(
a_panel, b_panel, c_panel, panel_m, k, b_n_group_stride, ldc);
} else {
neon_smmla::gemm_micro_smmla_8x8_packed_a(a_panel, b_panel, c_panel,
panel_m, k, ldc);
neon_smmla::gemm_micro_smmla_8x8_packed_a(
a_panel, b_panel + b_n_group_stride, c_panel + Nr, panel_m, k,
ldc);
}
}
}
}
};
} // namespace cpu_micro_gemm
#endif
csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp
+14 -8
View File
@@ -1,3 +1,6 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#ifndef CPU_MICRO_GEMM_NEON_HPP
#define CPU_MICRO_GEMM_NEON_HPP
@@ -16,9 +19,6 @@ namespace {
constexpr int32_t K = 4;
constexpr int32_t Cols = 2;
constexpr int32_t TileSize = K * Cols;
constexpr int32_t Mr = 8;
constexpr int32_t Nr = 8;
constexpr int32_t Nr_gemv = 16;
// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1]
FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) {
@@ -132,7 +132,7 @@ FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);
a_tile += 4 * TileSize;
b_tile += Nr * K;
b_tile += 4 * TileSize;
}
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
@@ -205,8 +205,8 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415);
a_tile += 2 * TileSize;
b_tile0 += Nr * K;
b_tile1 += Nr * K;
b_tile0 += 4 * TileSize;
b_tile1 += 4 * TileSize;
}
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
@@ -223,6 +223,9 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
template <typename scalar_t>
class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
public:
static constexpr int32_t Mr = 8;
static constexpr int32_t Nr = 8;
static constexpr int32_t NrGemv = 16;
static constexpr int32_t MaxMSize = 8;
static constexpr int32_t NSize = 32;
static constexpr int32_t WeightOCGroupSize = Nr;
@@ -246,6 +249,9 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
public:
using scalar_t = c10::BFloat16;
static constexpr int32_t Mr = 8;
static constexpr int32_t Nr = 8;
static constexpr int32_t NrGemv = 16;
static constexpr int32_t MaxMSize = 8;
static constexpr int32_t NSize = 32;
static constexpr int32_t WeightOCGroupSize = Nr;
@@ -253,7 +259,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
public:
// physical layout [
// M / 8; Mr is 8
// M / (8 or 4); Mr is 8 or 4
// K / 4; K for bfmmla is 4
// 4, ; 4 row-pairs for each 8 rows
// 2, ; row-pair is 2 rows
@@ -439,7 +445,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
(void)lda; // A is packed, so lda is not needed
TORCH_CHECK_EQ(k % K, 0);
for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
const bfloat16_t* __restrict__ b_panel =
reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
csrc/cpu/sgl-kernels/conv.cpp
+173 -12
View File
@@ -451,6 +451,90 @@ void causal_conv1d_update_kernel_impl(
});
}
template <typename scalar_t>
void causal_conv1d_update_multi_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
scalar_t* __restrict__ conv_states,
const scalar_t* __restrict__ weight,
const scalar_t* __restrict__ bias,
const int32_t* __restrict__ num_accepted_tokens,
const int32_t* __restrict__ conv_indices,
bool silu_activation,
int64_t batch,
int64_t dim,
int64_t seqlen,
int64_t width,
int64_t state_len,
int64_t conv_state_slot_stride) {
constexpr int64_t BLOCK_N = block_size_n() * 2;
const int64_t NB = div_up(dim, BLOCK_N);
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
int64_t bs{0}, nb{0};
data_index_init(begin, bs, batch, nb, NB);
for (int64_t i = begin; i < end; ++i) {
const int64_t nb_start = nb * BLOCK_N;
const int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
const int32_t conv_state_index = conv_indices[bs];
const int32_t history_offset = num_accepted_tokens[bs] - 1;
switch (width << 4 | nb_size >> 4) {
case 0x42:
tinygemm_kernel<scalar_t, 4, 32, has_bias, has_silu>::apply(
input + bs * seqlen * dim + nb_start,
weight + nb_start * width,
out + bs * seqlen * dim + nb_start,
has_bias ? bias + nb_start : nullptr,
conv_states + conv_state_index * conv_state_slot_stride +
history_offset * dim + nb_start,
true,
seqlen,
dim,
true);
break;
case 0x44:
tinygemm_kernel<scalar_t, 4, 64, has_bias, has_silu>::apply(
input + bs * seqlen * dim + nb_start,
weight + nb_start * width,
out + bs * seqlen * dim + nb_start,
has_bias ? bias + nb_start : nullptr,
conv_states + conv_state_index * conv_state_slot_stride +
history_offset * dim + nb_start,
true,
seqlen,
dim,
true);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
}
data_index_step(bs, batch, nb, NB);
}
});
});
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t bs = begin; bs < end; ++bs) {
const int32_t conv_state_index = conv_indices[bs];
const int32_t num_accepted = num_accepted_tokens[bs];
scalar_t* state = conv_states + conv_state_index * conv_state_slot_stride;
std::memmove(
state,
state + num_accepted * dim,
(state_len - seqlen) * dim * sizeof(scalar_t));
std::memcpy(
state + (state_len - seqlen) * dim,
input + bs * seqlen * dim,
seqlen * dim * sizeof(scalar_t));
}
});
}
} // anonymous namespace
// from [dim, width] or [N, K]
@@ -545,7 +629,7 @@ at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t n
// query_start_loc: (batch + 1) int32
// cache_indices: (batch) int32
// has_initial_state: (batch) bool
// conv_states: (..., dim, width - 1) itype
// conv_states: (..., dim, state_len) itype, where state_len >= width - 1
// activation: either None or "silu" or "swish"
// pad_slot_id: int
//
@@ -586,11 +670,14 @@ at::Tensor causal_conv1d_fwd_cpu(
CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
CHECK_GE(padded_batch, batch);
CHECK_EQ(conv_states_val.size(1), dim);
CHECK_EQ(conv_states_val.size(2), width - 1);
const int64_t state_len = conv_states_val.size(2);
CHECK_GE(state_len, width - 1);
// adjust `conv_states` to be contiguous on `dim`
// should happen only once
if (conv_states_val.stride(-2) != 1) {
TORCH_CHECK(state_len == width - 1,
"causal_conv1d_fwd_cpu: wide conv_states must be contiguous on dim.");
auto conv_states_copy = conv_states_val.clone();
conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
conv_states_val.copy_(conv_states_copy);
@@ -651,14 +738,14 @@ at::Tensor causal_conv1d_fwd_cpu(
// API aligned with GPUs
//
// x: (batch, dim) or (batch, dim, seqlen)
// x: (batch, dim) or (batch, seqlen, dim)
// conv_state: (..., dim, state_len), where state_len >= width - 1
// weight: (dim, width)
// bias: (dim,)
// cache_seqlens: (batch,), dtype int32.
// num_accepted_tokens: (batch,), dtype int32.
// conv_state_indices: (batch,), dtype int32
// pad_slot_id: int
// out: (batch, dim) or (batch, dim, seqlen)
// out: (batch, dim) or (batch, seqlen, dim)
//
at::Tensor causal_conv1d_update_cpu(
const at::Tensor& x,
@@ -666,7 +753,7 @@ at::Tensor causal_conv1d_update_cpu(
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
bool silu_activation,
const std::optional<at::Tensor>& cache_seqlens,
const std::optional<at::Tensor>& num_accepted_tokens,
const std::optional<at::Tensor>& conv_state_indices,
int64_t pad_slot_id,
bool is_vnni) {
@@ -674,13 +761,13 @@ at::Tensor causal_conv1d_update_cpu(
CHECK_CONTIGUOUS(weight);
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
// TODO: add multi-token prediction support
TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");
TORCH_CHECK(
x.dim() == 2 || x.dim() == 3,
"causal_conv1d_update_cpu: expect x to be 2D or 3D tensor.");
int64_t batch = x.size(0);
int64_t dim = x.size(1);
int64_t seqlen = 1;
int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2);
int64_t seqlen = x.dim() == 2 ? 1 : x.size(1);
int64_t width = weight.size(-1);
const auto scalar_type = x.scalar_type();
@@ -690,10 +777,84 @@ at::Tensor causal_conv1d_update_cpu(
CHECK_EQ(conv_states.scalar_type(), scalar_type);
CHECK_EQ(conv_states.size(1), dim);
CHECK_EQ(conv_states.size(2), width - 1);
const int64_t state_len = conv_states.size(2);
CHECK_GE(state_len, width - 1);
if (x.dim() == 3) {
TORCH_CHECK(
num_accepted_tokens.has_value(),
"causal_conv1d_update_cpu: num_accepted_tokens is required for 3D x.");
TORCH_CHECK(
conv_state_indices.has_value(),
"causal_conv1d_update_cpu: conv_state_indices is required for 3D x.");
CHECK_OPTIONAL_SHAPE_DTYPE(num_accepted_tokens, batch, at::kInt);
TORCH_CHECK(
width == 4,
"causal_conv1d_update_cpu: support only width of 4 for 3D x.");
TORCH_CHECK(
seqlen > 0,
"causal_conv1d_update_cpu: expect non-empty sequence for 3D x.");
TORCH_CHECK(
state_len >= seqlen,
"causal_conv1d_update_cpu: state_len must be >= seqlen for 3D x.");
TORCH_CHECK(
conv_states.stride(-2) == 1 && conv_states.stride(-1) == dim,
"causal_conv1d_update_cpu: 3D x requires SD conv_states layout.");
const int32_t* accepted_counts =
num_accepted_tokens.value().data_ptr<int32_t>();
const int32_t* indices = conv_state_indices.value().data_ptr<int32_t>();
const int64_t num_slots = conv_states.size(0);
for (int64_t bs = 0; bs < batch; ++bs) {
const int32_t num_accepted = accepted_counts[bs];
const int32_t conv_state_index = indices[bs];
TORCH_CHECK(
conv_state_index != pad_slot_id,
"causal_conv1d_update_cpu: 3D x does not support pad slots.");
TORCH_CHECK(
conv_state_index >= 0 && conv_state_index < num_slots,
"causal_conv1d_update_cpu: conv_state_indices out of range.");
TORCH_CHECK(
num_accepted >= 1 && num_accepted <= seqlen,
"causal_conv1d_update_cpu: num_accepted_tokens must be in [1, "
"seqlen].");
TORCH_CHECK(
num_accepted - 1 + width - 1 <= state_len,
"causal_conv1d_update_cpu: history window exceeds conv_states.");
}
int64_t conv_state_slot_stride = conv_states.stride(0);
at::Tensor out = at::empty_like(x);
AT_DISPATCH_REDUCED_FLOATING_TYPES(
scalar_type, "causal_conv1d_update_multi_kernel_impl", [&] {
causal_conv1d_update_multi_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
x.data_ptr<scalar_t>(),
conv_states.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
conditional_data_ptr<scalar_t>(bias),
accepted_counts,
indices,
silu_activation,
batch,
dim,
seqlen,
width,
state_len,
conv_state_slot_stride);
});
return out;
}
TORCH_CHECK(
!num_accepted_tokens.has_value(),
"causal_conv1d_update_cpu: num_accepted_tokens is only supported for 3D "
"x.");
// adjust `conv_states` to be contiguous on `dim`
if (conv_states.stride(-2) != 1) {
TORCH_CHECK(state_len == width - 1,
"causal_conv1d_update_cpu: wide conv_states must be contiguous on dim.");
int64_t num_cache_lines = conv_states.size(0);
auto conv_states_copy = conv_states.clone();
conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
csrc/cpu/torch_bindings.cpp
+34 -5
View File
@@ -147,7 +147,7 @@ at::Tensor causal_conv1d_fwd_cpu(
at::Tensor causal_conv1d_update_cpu(
const at::Tensor& x, const at::Tensor& conv_states,
const at::Tensor& weight, const std::optional<at::Tensor>& bias,
bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens,
const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
bool is_vnni);
@@ -207,6 +207,20 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
const torch::Tensor& topk_id, const bool skip_weighted,
const std::string& act, const std::string& isa);
void prepack_moe_weight_int8(const torch::Tensor& weight,
torch::Tensor& packed_weight,
const std::string& isa);
void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
const torch::Tensor& w13, const torch::Tensor& w2,
const torch::Tensor& w13_scale,
const torch::Tensor& w2_scale,
const std::optional<torch::Tensor>& w13_bias,
const std::optional<torch::Tensor>& w2_bias,
const torch::Tensor& topk_weights,
const torch::Tensor& topk_id, const bool skip_weighted,
const std::string& act, const std::string& isa);
void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
const torch::Tensor positions,
const torch::Tensor block_table,
@@ -502,7 +516,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "
"pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
@@ -596,8 +611,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#endif
// fused moe
#if defined(__AVX512F__) || \
(defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
ops.def(
"prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
"-> ()");
@@ -608,7 +622,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"bool skip_weighted, "
"str act, str isa) -> ()");
ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
#endif
#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) &&
// !defined(__APPLE__))
#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
!defined(__APPLE__)
ops.def(
"prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, "
"str isa) -> ()");
ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8);
ops.def(
"cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, "
"Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, "
"Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool "
"skip_weighted, str act, str isa) -> ()");
ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8);
#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) &&
// !defined(__APPLE__)
ops.def(
"mla_decode_kvcache("
" Tensor! out, Tensor query, Tensor kv_cache,"
csrc/fs_io.cpp
+283 -1
View File
@@ -3,19 +3,155 @@
#include <Python.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <filesystem>
#include <string>
#include <vector>
#if defined(O_DIRECT)
constexpr int kODirectFlag = O_DIRECT;
#else
constexpr int kODirectFlag = 0;
#endif
extern "C" {
static void _batch_lookup(const std::vector<const char*>& paths,
namespace {
// Returns 0 on success, or the std::error_code's POSIX-compatible value on
// failure, mirroring the errno convention used by the syscalls below.
inline int ensure_parent_dirs(const std::string& path) {
const auto parent = std::filesystem::path(path).parent_path();
if (parent.empty()) {
return 0;
}
std::error_code ec;
std::filesystem::create_directories(parent, ec);
return ec ? ec.value() : 0;
}
// Core single-block store: src/size are raw pointer + byte count. Returns 0
// on success, or the errno of the failing step on failure -- captured
// before any subsequent cleanup call can overwrite it. On failure, the temp
// file is removed.
inline int _store_block(const char* tmp_path, const char* dest_path,
const char* src, size_t size, bool use_o_direct) {
if (access(dest_path, F_OK) == 0) {
return 0; // Already present.
}
if (const int err = ensure_parent_dirs(dest_path); err != 0) {
return err;
}
const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
const int fd = open(
tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644);
if (fd < 0) {
return errno;
}
const ssize_t written = write(fd, src, size);
if (written < 0 || static_cast<size_t>(written) != size) {
const int err = written < 0 ? errno : EIO;
close(fd); // Best-effort cleanup; the real error is already captured.
unlink(tmp_path);
return err;
}
if (close(fd) != 0) {
const int err = errno;
unlink(tmp_path);
return err;
}
if (rename(tmp_path, dest_path) != 0) {
const int err = errno;
unlink(tmp_path);
return err;
}
return 0;
}
// Core single-block load: dst/size are raw pointer + byte count. Returns 0
// on success, or the errno of the failing step on failure. On failure,
// the source file is removed since a partially-read block should not be reused.
inline int _load_block(const char* source_path, char* dst, size_t size,
bool use_o_direct) {
const int o_direct_flag = use_o_direct ? kODirectFlag : 0;
const int fd = open(source_path, O_RDONLY | o_direct_flag, 0);
if (fd < 0) {
const int err = errno;
unlink(source_path);
return err;
}
const ssize_t bytes_read = read(fd, dst, size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != size) {
const int err = bytes_read < 0 ? errno : EIO;
close(fd);
unlink(source_path);
return err;
}
if (close(fd) != 0) {
const int err = errno;
unlink(source_path);
return err;
}
return 0;
}
inline void _batch_lookup(const std::vector<const char*>& paths,
std::vector<int>& exists_flags) {
for (size_t i = 0; i < paths.size(); i++) {
exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0;
}
}
// Helper: extract a list[str] of length n into a vector<const char*>.
// Returns false and sets a Python exception on error.
inline bool extract_str_list(PyObject* list, Py_ssize_t n,
std::vector<const char*>& out) {
for (Py_ssize_t i = 0; i < n; i++) {
out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr);
if (out[i] == nullptr) {
return false;
}
}
return true;
}
// Helper: extract a Py_buffer per element of a list[bytes-like] of length n.
// On success, `out` holds n acquired buffers (caller must PyBuffer_Release
// each). On failure, any buffers already acquired are released before
// returning false, and a Python exception is set.
inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags,
std::vector<Py_buffer>& out) {
for (Py_ssize_t i = 0; i < n; i++) {
if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) {
for (Py_ssize_t j = 0; j < i; j++) {
PyBuffer_Release(&out[j]);
}
return false;
}
}
return true;
}
inline void release_buffer_list(std::vector<Py_buffer>& buffers) {
for (auto& buf : buffers) {
PyBuffer_Release(&buf);
}
}
} // namespace
/// @brief Check file existence for a batch of paths.
/// @param paths list[str] – absolute paths to check.
/// @return list[bool] – True if the corresponding path exists, False otherwise.
@@ -51,11 +187,157 @@ static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) {
return result;
}
/// @brief Store a batch of blocks, each from its own buffer, to disk.
/// @param tmp_paths list[str] – one temp path per block.
/// @param dest_paths list[str] – one destination path per block.
/// @param buffers list[bytes-like] – one source buffer per block.
/// @param use_o_direct bool – whether to open files with O_DIRECT
/// (default True). Ignored where O_DIRECT is unsupported
/// by the platform.
/// @note Releases the GIL for the entire batch. Raises on first error.
static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) {
PyObject* tmp_paths_obj = nullptr;
PyObject* dest_paths_obj = nullptr;
PyObject* buffers_obj = nullptr;
int use_o_direct = 1;
if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj,
&PyList_Type, &dest_paths_obj, &PyList_Type,
&buffers_obj, &use_o_direct)) {
return nullptr;
}
const Py_ssize_t n = PyList_Size(tmp_paths_obj);
if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) {
PyErr_SetString(
PyExc_ValueError,
"tmp_paths, dest_paths and buffers must have the same length");
return nullptr;
}
std::vector<const char*> tmp_paths(n);
std::vector<const char*> dest_paths(n);
if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr;
if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr;
std::vector<Py_buffer> buffers(n);
if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) {
return nullptr;
}
Py_ssize_t failed_index = -1;
int failure_errno = 0;
{
Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
const char* buf = static_cast<const char*>(buffers[i].buf);
const int err =
_store_block(tmp_paths[i], dest_paths[i], buf,
static_cast<size_t>(buffers[i].len), use_o_direct);
if (err != 0) {
failed_index = i;
failure_errno = err;
break;
}
}
Py_END_ALLOW_THREADS
}
release_buffer_list(buffers);
if (failed_index >= 0) {
// PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
errno = failure_errno;
return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
dest_paths[failed_index]);
}
Py_RETURN_NONE;
}
/// @brief Load a batch of blocks from disk, each into its own buffer.
/// @param source_paths list[str] – one source path per block.
/// @param buffers list[writable bytes-like] – one destination buffer
/// per block.
/// @param use_o_direct bool – whether to open files with O_DIRECT
/// (default True). Ignored where O_DIRECT is unsupported
/// by the platform.
/// @note Releases the GIL for the entire batch. Raises on first error.
static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) {
PyObject* source_paths_obj = nullptr;
PyObject* buffers_obj = nullptr;
int use_o_direct = 1;
if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj,
&PyList_Type, &buffers_obj, &use_o_direct)) {
return nullptr;
}
const Py_ssize_t n = PyList_Size(source_paths_obj);
if (PyList_Size(buffers_obj) != n) {
PyErr_SetString(PyExc_ValueError,
"source_paths and buffers must have the same length");
return nullptr;
}
std::vector<const char*> source_paths(n);
if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr;
std::vector<Py_buffer> buffers(n);
if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) {
return nullptr;
}
Py_ssize_t failed_index = -1;
int failure_errno = 0;
{
Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) {
char* buf = static_cast<char*>(buffers[i].buf);
const int err =
_load_block(source_paths[i], buf, static_cast<size_t>(buffers[i].len),
use_o_direct);
if (err != 0) {
failed_index = i;
failure_errno = err;
break;
}
}
Py_END_ALLOW_THREADS
}
release_buffer_list(buffers);
if (failed_index >= 0) {
// PyErr_SetFromErrnoWithFilename() reads the errno to format exception.
errno = failure_errno;
return PyErr_SetFromErrnoWithFilename(PyExc_OSError,
source_paths[failed_index]);
}
Py_RETURN_NONE;
}
static PyMethodDef fs_io_C_methods[] = {
{"batch_lookup", batch_lookup, METH_VARARGS,
"batch_lookup(paths: list[str]) -> list[bool]\n"
"\n"
"Check file existence for a batch of paths."},
{"batch_store_block", batch_store_block, METH_VARARGS,
"batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n"
" buffers: list[bytes-like],\n"
" use_o_direct: bool = True) -> None\n"
"\n"
"Store a batch of blocks, each from its own buffer, to disk. Raises on "
"first error."},
{"batch_load_block", batch_load_block, METH_VARARGS,
"batch_load_block(source_paths: list[str],\n"
" buffers: list[writable bytes-like],\n"
" use_o_direct: bool = True) -> None\n"
"\n"
"Load a batch of blocks from disk into corresponding buffers. "
"Raises on first error."},
{nullptr, nullptr, 0, nullptr},
};
csrc/libtorch_stable/bf16_skinny_gemm.cu
-262
View File
@@ -1,262 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Skinny GEMM: activation(bf16) x weight(bf16)^T -> bf16, for decode-time
// M <= 32 with a large reduction dim. Replaces cuBLAS splitK (GEMM +
// splitKreduce) with a single block-per-output-column kernel; these shapes
// are weight-bandwidth-bound, so one coalesced pass over the weight at
// fp32 accumulation is optimal. Adapted from fp32_router_gemm.cu.
//
// First user: the DeepSeek-V32/GLM-5.2 MTP eh_proj (K=2*hidden=12288,
// N=hidden/TP), whose cuBLAS splitK pick costs ~34.6us vs the ~19us
// bandwidth floor per replicated read (and ~4us once column-parallel).
#include <cuda_bf16.h>
#include <cuda_runtime.h>
// ---------------------------------------------------------------------------
// Load helpers (8 x bf16 = one uint4 load, converted to fp32)
// ---------------------------------------------------------------------------
namespace skinny {
constexpr int VPT = 8; // bf16 values per thread per load
__device__ __forceinline__ void load_bf16x8(__nv_bfloat16 const* ptr,
float* dst) {
uint4 v = *reinterpret_cast<uint4 const*>(ptr);
__nv_bfloat16 const* p = reinterpret_cast<__nv_bfloat16 const*>(&v);
#pragma unroll
for (int i = 0; i < VPT; i++) dst[i] = __bfloat162float(p[i]);
}
// Streaming variant for the weight: each row is read exactly once across the
// whole grid, so bypass L2 residency (evict-first). Measured -1.4us at M=1.
__device__ __forceinline__ void load_bf16x8_cs(__nv_bfloat16 const* ptr,
float* dst) {
uint4 v = __ldcs(reinterpret_cast<uint4 const*>(ptr));
__nv_bfloat16 const* p = reinterpret_cast<__nv_bfloat16 const*>(&v);
#pragma unroll
for (int i = 0; i < VPT; i++) dst[i] = __bfloat162float(p[i]);
}
// ---------------------------------------------------------------------------
// Kernel: each block computes kNPB output columns for all kNumTokens rows.
// grid = kN / kNPB, block = kBlockSize threads. K is reduced VPT elements
// per thread per iteration; fp32 accumulation, warp butterfly + smem
// finalize, bf16 store.
// ---------------------------------------------------------------------------
template <int kBlockSize, int kNumTokens, int kNPB, int kPF, int kN, int kK>
__global__ __launch_bounds__(kBlockSize, 1) void bf16_skinny_gemm_kernel(
__nv_bfloat16* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
int64_t out_stride) {
constexpr int k_elems_per_iter = VPT * kBlockSize;
constexpr int k_iterations = kK / k_elems_per_iter;
static_assert(kK % k_elems_per_iter == 0);
constexpr int kWarpSize = 32;
constexpr int kNumWarps = kBlockSize / kWarpSize;
int const n_base = blockIdx.x * kNPB;
int const tid = threadIdx.x;
int const warpId = tid / kWarpSize;
int const laneId = tid % kWarpSize;
float acc[kNumTokens][kNPB] = {};
__shared__ float sm_reduction[kNumTokens][kNPB][kNumWarps];
// Register prefetch (kPF > 0): W does not depend on the predecessor, so
// the first kPF iterations' weight chunks are loaded raw BEFORE the
// dependency sync; with a PDL-releasing producer (fused_eh_norm fires
// gdc_launch_dependents early) these DRAM round trips overlap the norm.
// Pair-measured on B300 (norm+gemm in one graph, full 6144x12288):
// M=1 pf2 28.43us vs pf0 29.00us; deeper prefetch or M >= 2 regresses
// (register pressure), hence the per-M selection in the launcher.
uint4 w_pre[kPF > 0 ? kPF : 1][kNPB];
#pragma unroll
for (int pf = 0; pf < kPF; pf++) {
#pragma unroll
for (int n = 0; n < kNPB; n++) {
w_pre[pf][n] =
*reinterpret_cast<uint4 const*>(mat_b + (size_t)(n_base + n) * kK +
pf * k_elems_per_iter + tid * VPT);
}
}
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
#pragma unroll
for (int ki = 0; ki < k_iterations; ki++) {
int const k_base = ki * k_elems_per_iter + tid * VPT;
float b_float[kNPB][VPT];
if (ki < kPF) {
#pragma unroll
for (int n = 0; n < kNPB; n++) {
__nv_bfloat16 const* p =
reinterpret_cast<__nv_bfloat16 const*>(&w_pre[ki][n]);
#pragma unroll
for (int v = 0; v < VPT; v++) b_float[n][v] = __bfloat162float(p[v]);
}
} else {
#pragma unroll
for (int n = 0; n < kNPB; n++) {
load_bf16x8_cs(mat_b + (size_t)(n_base + n) * kK + k_base, b_float[n]);
}
}
#pragma unroll
for (int m = 0; m < kNumTokens; m++) {
float a_float[VPT];
load_bf16x8(mat_a + (size_t)m * kK + k_base, a_float);
#pragma unroll
for (int n = 0; n < kNPB; n++) {
#pragma unroll
for (int k = 0; k < VPT; k++) {
acc[m][n] += a_float[k] * b_float[n][k];
}
}
}
}
// Warp-level butterfly reduction
#pragma unroll
for (int m = 0; m < kNumTokens; m++) {
#pragma unroll
for (int n = 0; n < kNPB; n++) {
float sum = acc[m][n];
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][n][warpId] = sum;
}
}
__syncthreads();
// Parallel finalize: one thread per (m, n) output.
for (int idx = tid; idx < kNumTokens * kNPB; idx += kBlockSize) {
int const m = idx / kNPB;
int const n = idx % kNPB;
float final_sum = 0.0f;
#pragma unroll
for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][n][w];
out[(size_t)m * out_stride + n_base + n] = __float2bfloat16(final_sum);
}
// Trigger after our stores: harmless hardening, not a guarantee — the
// trigger only permits dependent-launch scheduling and carries no memory
// visibility semantics, so a PDL consumer must still gridsync before
// reading our output. Every current consumer is a plain launch (full
// stream order); firing late just avoids an unnecessarily early launch
// window and matches fp32_router_gemm.cu.
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
}
} // namespace skinny
// ---------------------------------------------------------------------------
// Launcher
// ---------------------------------------------------------------------------
template <int kBlockSize, int kNPB, int kNumTokens, int kN, int kK>
void invokeBf16SkinnyGemm(__nv_bfloat16* output, __nv_bfloat16 const* mat_a,
__nv_bfloat16 const* mat_b, int64_t out_stride,
cudaStream_t stream) {
static_assert(kN % kNPB == 0);
// Weight prefetch depth: only M=1 measured a win (see kernel comment).
constexpr int kPF = (kNumTokens == 1) ? 2 : 0;
cudaLaunchConfig_t config;
config.gridDim = kN / kNPB;
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,
skinny::bf16_skinny_gemm_kernel<kBlockSize, kNumTokens,
kNPB, kPF, kN, kK>,
output, mat_a, mat_b, out_stride);
}
// ---------------------------------------------------------------------------
// Explicit instantiations. M = 1..32; (N, K) pairs:
// (768, 12288) eh_proj shard, TP8
// (1536, 12288) eh_proj shard, TP4
// (6144, 12288) eh_proj unsharded
// kNPB (B300 sweep, M=1): 6144 -> 2 (3072 blocks, 23.2us vs 25.3 at kNPB=8;
// narrow blocks minimize wave quantization); shards 768/1536 keep 4.
// ---------------------------------------------------------------------------
#define INSTANTIATE(M, NPB, N, K) \
template void invokeBf16SkinnyGemm<128, NPB, M, N, K>( \
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int64_t, \
cudaStream_t);
#define INSTANTIATE_ALL_M(NPB, N, K) \
INSTANTIATE(1, NPB, N, K) \
INSTANTIATE(2, NPB, N, K) \
INSTANTIATE(3, NPB, N, K) \
INSTANTIATE(4, NPB, N, K) \
INSTANTIATE(5, NPB, N, K) \
INSTANTIATE(6, NPB, N, K) \
INSTANTIATE(7, NPB, N, K) \
INSTANTIATE(8, NPB, N, K) \
INSTANTIATE(9, NPB, N, K) \
INSTANTIATE(10, NPB, N, K) \
INSTANTIATE(11, NPB, N, K) \
INSTANTIATE(12, NPB, N, K) \
INSTANTIATE(13, NPB, N, K) \
INSTANTIATE(14, NPB, N, K) \
INSTANTIATE(15, NPB, N, K) \
INSTANTIATE(16, NPB, N, K) \
INSTANTIATE(17, NPB, N, K) \
INSTANTIATE(18, NPB, N, K) \
INSTANTIATE(19, NPB, N, K) \
INSTANTIATE(20, NPB, N, K) \
INSTANTIATE(21, NPB, N, K) \
INSTANTIATE(22, NPB, N, K) \
INSTANTIATE(23, NPB, N, K) \
INSTANTIATE(24, NPB, N, K) \
INSTANTIATE(25, NPB, N, K) \
INSTANTIATE(26, NPB, N, K) \
INSTANTIATE(27, NPB, N, K) \
INSTANTIATE(28, NPB, N, K) \
INSTANTIATE(29, NPB, N, K) \
INSTANTIATE(30, NPB, N, K) \
INSTANTIATE(31, NPB, N, K) \
INSTANTIATE(32, NPB, N, K)
INSTANTIATE_ALL_M(4, 768, 12288)
INSTANTIATE_ALL_M(4, 1536, 12288)
INSTANTIATE_ALL_M(2, 6144, 12288)
// LL-mode (M<=8 wiring guard) backbone shapes, B300 sweep vs cuBLAS:
// q_b_proj (2048, 2048): 1.67x/1.29x/1.15x at M=4/6/8 (NPB=4 within
// 0.1us of per-M best)
// shared-expert gate_up (512, 6144): 1.95x/1.58x/1.40x at M=4/6/8
// cuBLAS keeps qkv_a (2624,6144) and o_proj (6144,2048) — already at
// 3.4-3.8 TB/s there; the GEMV loses on activation re-reads.
INSTANTIATE_ALL_M(4, 2048, 2048)
INSTANTIATE_ALL_M(4, 512, 6144)
// fused_qkv_a (2624, 6144), 32MB: skinny wins ONLY at M<=2 (B300: M=1
// 6.99us vs cuBLAS 9.21 = 1.32x, M=2 1.24x; M>=4 cuBLAS holds at 3.5TB/s
// and every alternative loses — cublasLt top-8 3.6TB/s wall, DeepGEMM
// 0.71x, wmma+cp.async custom 0.30x pending a TMA rewrite).
INSTANTIATE_ALL_M(4, 2624, 6144)
// DSv3.2 (TP8) siblings of the GLM shapes above, same dual-chip matrix:
// fused_qkv_a (2112, 7168), 30MB: wins M<=2 (M=1 1.30-1.34x)
// MTP eh_proj (7168, 14336), 205MB: wins M<=2 (M=1 1.12-1.16x)
INSTANTIATE_ALL_M(4, 2112, 7168)
INSTANTIATE_ALL_M(2, 7168, 14336)
#undef INSTANTIATE_ALL_M
#undef INSTANTIATE
csrc/libtorch_stable/bf16_skinny_gemm_entry.cu
-170
View File
@@ -1,170 +0,0 @@
// 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 SKINNY_MAX_TOKENS = 32;
// Supported (N, K) pairs (must match the instantiations in
// bf16_skinny_gemm.cu): eh_proj shard TP8 / TP4 / unsharded.
static inline bool bf16_skinny_gemm_supported(int n, int k) {
if (k == 12288 && (n == 768 || n == 1536 || n == 6144)) return true;
// LL-mode backbone shapes (wire callers with an M <= 8 guard; the GEMV
// family loses to cuBLAS at larger M).
if (k == 2048 && n == 2048) return true; // q_b_proj (TP8)
if (k == 6144 && n == 2624) return true; // fused_qkv_a (wire M <= 2 only)
if (k == 7168 && n == 2112) return true; // DSv3.2 fused_qkv_a (M <= 2)
if (k == 14336 && n == 7168) return true; // DSv3.2 eh_proj (M <= 2)
if (k == 6144 && n == 512) return true; // shared-expert gate_up (TP8)
return false;
}
// Forward declarations - template params must match bf16_skinny_gemm.cu
template <int kBlockSize, int kNPB, int kNumTokens, int kN, int kK>
void invokeBf16SkinnyGemm(__nv_bfloat16* output, __nv_bfloat16 const* mat_a,
__nv_bfloat16 const* mat_b, int64_t out_stride,
cudaStream_t stream);
template <int kNPB, int kN, int kK, int kBegin, int kEnd>
struct SkinnyLoopUnroller {
static void unroll(int num_tokens, __nv_bfloat16* output,
__nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
int64_t out_stride, cudaStream_t stream) {
if (num_tokens == kBegin) {
invokeBf16SkinnyGemm<128, kNPB, kBegin, kN, kK>(output, mat_a, mat_b,
out_stride, stream);
} else {
SkinnyLoopUnroller<kNPB, kN, kK, kBegin + 1, kEnd>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
}
}
};
template <int kNPB, int kN, int kK, int kEnd>
struct SkinnyLoopUnroller<kNPB, kN, kK, kEnd, kEnd> {
static void unroll(int num_tokens, __nv_bfloat16* output,
__nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
int64_t out_stride, cudaStream_t stream) {
if (num_tokens == kEnd) {
invokeBf16SkinnyGemm<128, kNPB, kEnd, kN, kK>(output, mat_a, mat_b,
out_stride, stream);
} else {
throw std::invalid_argument(
"bf16_skinny_gemm: num_tokens must be in [1, 32]");
}
}
};
static void dispatchBf16SkinnyGemm(int n, int k, int num_tokens,
__nv_bfloat16* output,
__nv_bfloat16 const* mat_a,
__nv_bfloat16 const* mat_b,
int64_t out_stride, cudaStream_t stream) {
if (n == 768 && k == 12288) {
SkinnyLoopUnroller<4, 768, 12288, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 1536 && k == 12288) {
SkinnyLoopUnroller<4, 1536, 12288, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 6144 && k == 12288) {
SkinnyLoopUnroller<2, 6144, 12288, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 2048 && k == 2048) {
SkinnyLoopUnroller<4, 2048, 2048, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 2624 && k == 6144) {
SkinnyLoopUnroller<4, 2624, 6144, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 2112 && k == 7168) {
SkinnyLoopUnroller<4, 2112, 7168, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 7168 && k == 14336) {
SkinnyLoopUnroller<2, 7168, 14336, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else if (n == 512 && k == 6144) {
SkinnyLoopUnroller<4, 512, 6144, 1, SKINNY_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, out_stride, stream);
} else {
throw std::invalid_argument("bf16_skinny_gemm: unsupported (N, K) pair");
}
}
void bf16_skinny_gemm(
torch::stable::Tensor& output, // [num_tokens, N] bf16
torch::stable::Tensor const& mat_a, // [num_tokens, K] bf16
torch::stable::Tensor const& mat_b // [N, K] bf16
) {
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(),
"bf16_skinny_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(),
"bf16_skinny_gemm: all tensors must be on the same device");
STD_TORCH_CHECK(mat_a.is_contiguous() && mat_b.is_contiguous(),
"bf16_skinny_gemm: inputs must be contiguous");
// Output may be a column-slice view of a wider padded buffer: unit column
// stride, row stride >= N (rows must not overlap).
STD_TORCH_CHECK(output.stride(1) == 1,
"bf16_skinny_gemm: output columns must be contiguous");
const int num_tokens = mat_a.size(0);
const int n = mat_b.size(0);
const int k = mat_a.size(1);
STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == n,
"bf16_skinny_gemm: output must be [num_tokens, N]");
STD_TORCH_CHECK(mat_b.size(1) == k,
"bf16_skinny_gemm: mat_a and mat_b must share K");
STD_TORCH_CHECK(bf16_skinny_gemm_supported(n, k),
"bf16_skinny_gemm: unsupported (N, K) pair");
const int64_t out_stride = output.stride(0);
STD_TORCH_CHECK(num_tokens <= 1 || out_stride >= n,
"bf16_skinny_gemm: output rows overlap");
STD_TORCH_CHECK(num_tokens >= 0 && num_tokens <= SKINNY_MAX_TOKENS,
"bf16_skinny_gemm: num_tokens must be in [0, 32]");
STD_TORCH_CHECK(
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"bf16_skinny_gemm: all tensors must be bfloat16");
// Empty batch (e.g. an empty rank at a DP/PP boundary): nothing to compute.
if (num_tokens == 0) {
return;
}
const torch::stable::accelerator::DeviceGuard device_guard(
mat_a.get_device_index());
STD_TORCH_CHECK(getSMVersion() >= 90, "bf16_skinny_gemm: requires SM90+");
auto stream = get_current_cuda_stream(mat_a.get_device_index());
dispatchBf16SkinnyGemm(
n, k, num_tokens,
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), out_stride,
stream);
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("bf16_skinny_gemm", TORCH_BOX(&bf16_skinny_gemm));
}
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
+36 -112
View File
@@ -391,7 +391,7 @@ struct MmaComputer {
static constexpr int n_iter_cnt =
(tile_n + 7) /
8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
static_assert(m_iter_cnt == 1 || m_iter_cnt == 2);
static_assert(m_iter_cnt == 1);
static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
__device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
@@ -416,18 +416,13 @@ struct MmaComputer {
public:
__device__ void prepare() {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
// Fragment addressing is per 16-row ldmatrix tile; m_iter selects the
// 16-row half within tile_m.
#pragma unroll
for (int m = 0; m < m_iter_cnt; m++) {
#pragma unroll
for (int i = 0; i < k_phase_cnt; i++) {
int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
int m_idx = linear_idx % 16 + m * 16;
int k_idx = linear_idx / 16 + warp_k_offset_in_tile_k;
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
a_smem_offsets[m][i] = m_idx * tile_k + k_idx;
}
for (int i = 0; i < k_phase_cnt; i++) {
int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
int m_idx = linear_idx % tile_m;
int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
}
#pragma unroll
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
@@ -451,14 +446,11 @@ struct MmaComputer {
wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
#pragma unroll
for (int m = 0; m < m_iter_cnt; m++) {
#pragma unroll
for (int i = 0; i < k_phase_cnt; i++) {
int smem_offset = a_smem_offsets[m][i];
bf16_t* smem_ptr_this_iter =
smem_a + stage_idx * tile_m * tile_k + smem_offset;
ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[m][i]));
}
for (int i = 0; i < k_phase_cnt; i++) {
int smem_offset = a_smem_offsets[0][i];
bf16_t* smem_ptr_this_iter =
smem_a + stage_idx * tile_m * tile_k + smem_offset;
ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
}
#pragma unroll
@@ -477,12 +469,9 @@ struct MmaComputer {
for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
#pragma unroll
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
#pragma unroll
for (int m = 0; m < m_iter_cnt; m++) {
hmma_16_8_16_f32acc_bf16ab(
acc_reg[m][n_iter_idx], a_reg[m][k_iter_idx],
b_reg[n_iter_idx][k_iter_idx], acc_reg[m][n_iter_idx]);
}
hmma_16_8_16_f32acc_bf16ab(
acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
}
}
::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
@@ -497,14 +486,14 @@ struct MmaComputer {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
// reorganize the acc_reg
constexpr int thread_m = 2 * m_iter_cnt;
constexpr int thread_m = 2;
constexpr int thread_n = 2 * n_iter_cnt;
constexpr int cta_mma_n = n_iter_cnt * 8;
float acc_reg_reorg[thread_m][thread_n];
for (int i = 0; i < thread_m; i++) {
for (int j = 0; j < thread_n; j++) {
acc_reg_reorg[i][j] = acc_reg[i / 2][j / 2][(j % 2) + (i % 2) * 2];
acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
}
}
@@ -525,8 +514,7 @@ struct MmaComputer {
for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
#pragma unroll
for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
int m_idx =
(lane_idx / 4) + (m_idx_thread % 2) * 8 + (m_idx_thread / 2) * 16;
int m_idx = (lane_idx / 4) + m_idx_thread * 8;
int n_idx =
((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
@@ -599,7 +587,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
static_assert(
tile_k == 128 || tile_k == 256 || tile_k == 512 ||
tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
static_assert(tile_m == 16 || tile_m == 32);
static_assert(tile_m == 16);
constexpr int g2s_vec_bytes = 16;
constexpr int a_elem_bytes = 2;
constexpr int b_elem_bytes = 2;
@@ -659,7 +647,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
#endif
}
template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileM = 16>
template <typename T, int kHdIn, int kHdOut, int kTileN>
void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
cudaStream_t const stream) {
constexpr int gemm_m = kHdOut; // 2112
@@ -667,7 +655,7 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
constexpr int gemm_k = kHdIn; // 7168
constexpr int batch_size = 1;
std::swap(mat_a, mat_b);
constexpr int tile_m = kTileM;
constexpr int tile_m = 16;
constexpr int tile_n = kTileN; // 8 or 16
constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
constexpr int max_stage_cnt =
@@ -714,44 +702,6 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
// GLM-5.2 fused_qkv_a (K=6144 -> N=2624). tile_m=32 so the grid is
// 2624/32 = 82 CTAs, a single wave on B300 (148 SMs); tile_m=16's 164 CTAs
// straddle two waves and drop to 2.9 TB/s vs 3.9-4.0 here. Beats cuBLAS
// at every decode M: 1.11-1.15x for M<=8 (tile_n=8), 1.12x at M=16
// (tile_n=16). The M<=2 dispatch still belongs to bf16_skinny_gemm
// (4.5 TB/s).
template void invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 8, 32>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
template void invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 16, 32>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
// GLM-5.2 q_b_proj TP8 (K=2048 -> N=2048). tile_m=16 keeps 128 CTAs (already
// a single wave) and measures ahead of tile_m=32 here. Beats cuBLAS
// 1.79-1.87x at M=3..8 (tile_n=8) and 1.19-1.30x at M=9..16 (tile_n=16);
// M<=2 belongs to bf16_skinny_gemm, M>=20 to cuBLAS.
template void invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 8>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
template void invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 16>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
// DSv3.2 q_b_proj TP8 (K=1536 -> N=3072). tile_m=32 keeps 96 CTAs (single
// wave; tile_m=16's 192 straddle two). Beats cuBLAS 1.6-1.8x at M=1..8 and
// 1.05-1.18x at M=9..16 on B300/B200; the whole 1..16 range dispatches here
// (no skinny tier: K=1536 does not fit the GEMV's 128x8 K-step).
template void invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 8, 32>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
template void invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 16, 32>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
void dsv3_fused_a_gemm(torch::stable::Tensor& output,
torch::stable::Tensor const& mat_a,
torch::stable::Tensor const& mat_b) {
@@ -760,15 +710,12 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
int const hd_in = mat_a.size(1);
int const hd_out = mat_b.size(1);
bool const is_dsv3 = hd_in == 7168 && hd_out == 2112;
bool const is_glm = hd_in == 6144 && hd_out == 2624;
bool const is_glm_qb = hd_in == 2048 && hd_out == 2048;
bool const is_ds_qb = hd_in == 1536 && hd_out == 3072;
constexpr int kHdIn = 7168;
constexpr int kHdOut = 2112;
STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
"required 1 <= mat_a.shape[0] <= 16");
STD_TORCH_CHECK(is_dsv3 || is_glm || is_glm_qb || is_ds_qb,
"supported (hd_in, hd_out): (7168, 2112), (6144, 2624), "
"(2048, 2048), (1536, 3072)");
STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
STD_TORCH_CHECK(output.size(0) == num_tokens,
"required output.shape[0] == mat_a.shape[0]");
STD_TORCH_CHECK(output.size(1) == hd_out,
@@ -791,41 +738,18 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
auto stream = get_current_cuda_stream(mat_a.get_device_index());
auto* out = reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
auto* a = reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
auto* b = reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
if (is_dsv3) {
if (num_tokens <= 8) {
invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(out, a, b, num_tokens,
stream);
} else {
invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(out, a, b, num_tokens,
stream);
}
} else if (is_glm) {
if (num_tokens <= 8) {
invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 8, 32>(out, a, b, num_tokens,
stream);
} else {
invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 16, 32>(out, a, b, num_tokens,
stream);
}
} else if (is_glm_qb) {
if (num_tokens <= 8) {
invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 8>(out, a, b, num_tokens,
stream);
} else {
invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 16>(out, a, b, num_tokens,
stream);
}
if (num_tokens <= 8) {
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
stream);
} else {
if (num_tokens <= 8) {
invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 8, 32>(out, a, b, num_tokens,
stream);
} else {
invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 16, 32>(out, a, b, num_tokens,
stream);
}
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
stream);
}
}
csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu
+954
View File
@@ -0,0 +1,954 @@
/*
* Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
*/
// Production AttnRes forward for Blackwell (SM100).
//
// Warp-specialized online softmax + residual + RMSNorm:
// - 1 producer warp issues cp.async.bulk row loads into shared memory.
// - 8 consumer warps compute reductions and output.
// - Q=res_weight*rms_weight remains in registers across persistent tokens.
// - V rows are converted once and cached as FP32 in TMEM between passes.
//
// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major
// block residual storage.
#include "../torch_utils.h"
#include <cfloat>
#include <cstdint>
#include <cstdio>
#include <cuda_runtime.h>
#include <type_traits>
using bf16_t = __nv_bfloat16;
namespace sm100 {
namespace fwd_prod_v2 {
constexpr int K_TILE = 1024;
constexpr int N_CHUNK_DEFAULT = 4;
constexpr int CHUNK_DEPTH = 2;
constexpr int BLK = 288; // 1 producer warp + 8 consumer warps
constexpr int CONSUMER_THREADS = BLK - 32; // 256
constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32;
constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups
constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS;
constexpr int FIRST_USER_NAMED_BARRIER = 8;
__device__ __forceinline__ const bf16_t* residual_addr(
const bf16_t* block_res, const bf16_t* layer_res, int source, int N,
int token, int block_stride_m, int block_stride_r, int H) {
if (source < N - 1) {
return block_res + static_cast<long long>(token) * block_stride_m +
source * block_stride_r;
}
return layer_res + static_cast<long long>(token) * H;
}
__device__ __forceinline__ uint32_t elect_one_sync() {
uint32_t pred = 0;
uint32_t laneid = 0;
asm volatile(
"{\n"
".reg .b32 %%rx;\n"
".reg .pred %%px;\n"
" elect.sync %%rx|%%px, %2;\n"
"@%%px mov.s32 %1, 1;\n"
" mov.s32 %0, %%rx;\n"
"}\n"
: "+r"(laneid), "+r"(pred)
: "r"(0xffffffff));
return pred;
}
__device__ __forceinline__ void mbarrier_init(uint64_t& barrier,
int thread_count) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr),
"r"(thread_count));
}
__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier,
uint32_t bytes) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(
barrier_addr),
"r"(bytes));
}
__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile(
"{\n"
".reg .pred p;\n"
"WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n"
"@p bra DONE;\n"
"bra WAIT;\n"
"DONE:\n"
"}\n" ::"r"(barrier_addr),
"r"(phase));
}
__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
uint32_t const barrier_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
asm volatile(
"{\n"
".reg .b64 state;\n"
"mbarrier.arrive.shared::cta.b64 state, [%0];\n"
"}\n" ::"r"(barrier_addr));
}
__device__ __forceinline__ void fence_mbarrier_init() {
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads,
uint32_t user_barrier_id) {
asm volatile(
"bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER),
"r"(num_threads)
: "memory");
}
__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) {
uint32_t const dst_addr =
static_cast<uint32_t>(__cvta_generic_to_shared(dst));
asm volatile(
"tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
dst_addr),
"r"(num_columns));
}
__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) {
asm volatile(
"tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr),
"r"(num_columns));
}
__device__ __forceinline__ void tmem_release_allocation_lock() {
asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
}
__device__ __forceinline__ void tmem_store_wait() {
asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
}
template <int N, typename T>
__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) {
uint32_t* values = reinterpret_cast<uint32_t*>(dst);
if constexpr (N == 8) {
asm volatile(
"tcgen05.ld.sync.aligned.32x32b.x8.b32"
"{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
: "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]),
"=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7])
: "r"(src_addr));
} else {
static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
asm volatile(
"tcgen05.ld.sync.aligned.32x32b.x4.b32"
"{%0, %1, %2, %3}, [%4];\n"
: "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3])
: "r"(src_addr));
}
}
template <int N, typename T>
__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) {
uint32_t* values = reinterpret_cast<uint32_t*>(src);
if constexpr (N == 8) {
asm volatile(
"tcgen05.st.sync.aligned.32x32b.x8.b32"
"[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]),
"r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]),
"r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr));
} else {
static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
asm volatile(
"tcgen05.st.sync.aligned.32x32b.x4.b32"
"[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]),
"r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr));
}
}
__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) {
float2 result;
asm volatile("add.rn.f32x2 %0, %1, %2;\n"
: "=l"(reinterpret_cast<uint64_t&>(result))
: "l"(reinterpret_cast<uint64_t const&>(a)),
"l"(reinterpret_cast<uint64_t const&>(b)));
return result;
}
__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) {
float2 result;
asm volatile("mul.f32x2 %0, %1, %2;\n"
: "=l"(reinterpret_cast<uint64_t&>(result))
: "l"(reinterpret_cast<uint64_t const&>(a)),
"l"(reinterpret_cast<uint64_t const&>(b)));
return result;
}
__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b,
const float2& c) {
float2 result;
asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
: "=l"(reinterpret_cast<uint64_t&>(result))
: "l"(reinterpret_cast<uint64_t const&>(a)),
"l"(reinterpret_cast<uint64_t const&>(b)),
"l"(reinterpret_cast<uint64_t const&>(c)));
return result;
}
template <int NC>
struct FwdSmemPlan {
alignas(16) uint64_t bar_ready[CHUNK_DEPTH];
alignas(16) uint64_t bar_consumed[CHUNK_DEPTH];
alignas(16) uint64_t bar_output_norm_ready;
alignas(16) float2 ws_stats[CONSUMER_WARPS][NC];
uint32_t tmem_base;
};
__device__ __forceinline__ void cp_async_bulk(void* smem_dst,
const void* gmem_src, int bytes,
uint64_t& mbar) {
uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst));
uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar));
asm volatile(
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], "
"[%1], %2, [%3];\n" ::"r"(s),
"l"(gmem_src), "r"(bytes), "r"(m)
: "memory");
}
template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
bool OUTPUT_NORM_IN_SMEM = false>
__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T,
int B, int block_stride_m, int block_stride_r, float rms_eps,
const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
constexpr float LOG2_E = 1.4426950408889634f;
constexpr int N_CHUNK = NC;
// The two-source specialization only consumes half of the TMEM columns.
constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256;
constexpr int NUM_BUFS = CHUNK_DEPTH * NC;
constexpr int NHT = H / K_TILE;
constexpr int SLICES_PER_GROUP =
(NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS;
constexpr int VEC = 8;
constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC;
constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC;
constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP;
static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC);
static_assert(H >= 4096 && H <= 8192);
static_assert(H % K_TILE == 0);
const int tid = threadIdx.x;
const int wid = tid >> 5;
const int lane = tid & 31;
const int TB = T * B;
const int num_ctas = gridDim.x;
const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
const int comp_wid = wid - 1;
const int comp_tid = tid - 32;
const int group = (comp_wid >= 4) ? 1 : 0;
const int ct_in_group =
(comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1;
const int k_local = ct_in_group * VEC;
constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t);
constexpr size_t DELTA_BYTES =
HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0;
constexpr size_t OUTPUT_NORM_BYTES =
OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0;
extern __shared__ __align__(16) char smem_raw[];
bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw); // [NUM_BUFS][H]
bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES);
bf16_t* output_norm_buf =
reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES);
FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>(
smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES);
auto slot_of = [](long long gci, int n) {
return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n;
};
auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); };
auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; };
auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* {
return delta_bufs + chunk_slot * H;
};
if (wid == 0 && elect_one_sync()) {
#pragma unroll
for (int i = 0; i < CHUNK_DEPTH; i++) {
mbarrier_init(plan.bar_ready[i], 1);
mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS);
}
if constexpr (OUTPUT_NORM_IN_SMEM) {
mbarrier_init(plan.bar_output_norm_ready, 1);
}
fence_mbarrier_init();
}
// gdc wait BEFORE tmem alloc
cudaGridDependencySynchronize();
if (wid == 1) {
tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base);
if constexpr (RELEASE_TMEM) {
tmem_release_allocation_lock();
}
}
__syncthreads();
if constexpr (OUTPUT_NORM_IN_SMEM) {
if (wid == 0 && elect_one_sync()) {
mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t));
cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t),
plan.bar_output_norm_ready);
}
}
const uint32_t my_v_tmem =
comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0;
float q_cache[ACC_PER_THREAD];
if (comp_tid >= 0) {
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
#pragma unroll
for (int j = 0; j < 4; j++) {
int h = h_base + j;
q_cache[si * VEC + j] =
__bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
#pragma unroll
for (int j = 0; j < VEC; j++) {
int h = h_base + j;
q_cache[si * VEC + j] =
__bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
}
}
}
if (wid == 0) {
if (elect_one_sync()) {
long long gci = 0;
for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
const int t = tb / B;
for (int ci = 0; ci < num_chunks; ci++, gci++) {
int ns = ci * N_CHUNK;
int an = min(N_CHUNK, N - ns);
int chunk_slot = (int)(gci % CHUNK_DEPTH);
int pc = phase_of(gci);
mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1);
int transaction_bytes = an * H * (int)sizeof(bf16_t);
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (prefix_n >= 0 && prefix_n < an) {
transaction_bytes += H * (int)sizeof(bf16_t);
}
}
mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes);
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
if (n >= an) continue;
int slot = slot_of(gci, n);
const bf16_t* src =
residual_addr(block_res, layer_res, ns + n, N, t,
block_stride_m, block_stride_r, H);
cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t),
plan.bar_ready[chunk_slot]);
}
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (prefix_n >= 0 && prefix_n < an) {
cp_async_bulk(delta_buf_ptr(chunk_slot),
delta + (long long)tb * H, H * sizeof(bf16_t),
plan.bar_ready[chunk_slot]);
}
}
}
}
}
} else {
float acc32[ACC_PER_THREAD] = {};
float eps_cache;
asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps));
long long gci = 0;
for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
float m_running = -FLT_MAX;
float s_running = 0.f;
#pragma unroll
for (int i = 0; i < ACC_PER_THREAD; i++) {
acc32[i] = 0.f;
}
for (int ci = 0; ci < num_chunks; ci++, gci++) {
int ns = ci * N_CHUNK;
int an = min(N_CHUNK, N - ns);
int chunk_slot = (int)(gci % CHUNK_DEPTH);
int pr = phase_of(gci);
mbarrier_wait(plan.bar_ready[chunk_slot], pr);
float2 sq_local[N_CHUNK] = {};
float2 dot_local[N_CHUNK] = {};
auto pass_A_body = [&](auto AN_TOK) {
constexpr int AN = decltype(AN_TOK)::value;
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base =
6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
const float* qv = &q_cache[si * VEC];
#pragma unroll
for (int n = 0; n < AN; n++) {
int slot = slot_of(gci, n);
int2 vp =
*reinterpret_cast<const int2*>(buf_ptr(slot) + h_base);
auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (n == prefix_n) {
const bf16_t* delta_ptr =
delta_buf_ptr(chunk_slot) + h_base;
#pragma unroll
for (int j = 0; j < 2; j++) {
auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
delta_ptr + 2 * j);
v2[j] = __hadd2(v2[j], delta2);
}
*reinterpret_cast<int2*>(layer_res + (long long)tb * H +
h_base) = vp;
}
}
float2 f[2] = {__bfloat1622float2(v2[0]),
__bfloat1622float2(v2[1])};
tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
sq_local[n] = float2_fma(f[0], f[0], sq_local[n]);
sq_local[n] = float2_fma(f[1], f[1], sq_local[n]);
dot_local[n] =
float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]);
dot_local[n] =
float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]);
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
const float* qv = &q_cache[si * VEC];
#pragma unroll
for (int n = 0; n < AN; n++) {
int slot = slot_of(gci, n);
int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base);
auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
if constexpr (HAS_DELTA) {
int prefix_n = N - 1 - ns;
if (n == prefix_n) {
const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base;
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
delta_ptr + 2 * j);
v2[j] = __hadd2(v2[j], delta2);
}
*reinterpret_cast<int4*>(layer_res + (long long)tb * H +
h_base) = vp;
}
}
float2 f[4] = {
__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]),
__bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])};
tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
sq_local[n] = float2_fma(f[j], f[j], sq_local[n]);
dot_local[n] = float2_fma(
f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]);
}
}
}
};
if constexpr (NC == 4) {
switch (an) {
case 4:
pass_A_body(std::integral_constant<int, 4>{});
break;
case 3:
pass_A_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_A_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_A_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else if constexpr (NC == 3) {
switch (an) {
case 3:
pass_A_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_A_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_A_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else {
static_assert(NC == 2);
switch (an) {
case 2:
pass_A_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_A_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
}
if (lane == 0) {
mbarrier_arrive(plan.bar_consumed[chunk_slot]);
}
tmem_store_wait();
float2 reduce_pair[N_CHUNK];
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y,
dot_local[n].x + dot_local[n].y);
}
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]);
packed = __shfl_xor_sync(0xffffffff, packed, offset);
float2 other = reinterpret_cast<float2&>(packed);
reduce_pair[n] = float2_add(reduce_pair[n], other);
}
}
if (lane == 0) {
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
plan.ws_stats[comp_wid][n] = reduce_pair[n];
}
}
named_barrier_sync(CONSUMER_THREADS, 0);
float local_rsig = 0.f;
float local_logit = 0.f;
int stat_n = lane / CONSUMER_WARPS;
int stat_w = lane % CONSUMER_WARPS;
float2 totals = {};
if (stat_n < N_CHUNK) {
totals = plan.ws_stats[stat_w][stat_n];
}
#pragma unroll
for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
totals.x +=
__shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS);
totals.y +=
__shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS);
}
if (stat_n < N_CHUNK && stat_w == 0) {
local_rsig = rsqrtf(totals.x / H + eps_cache);
local_logit = totals.y * local_rsig;
}
float logit_n[N_CHUNK];
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS);
}
float m_chunk = -FLT_MAX;
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]);
}
float m_new = fmaxf(m_running, m_chunk);
float corr = exp2f((m_running - m_new) * LOG2_E);
float w_n[N_CHUNK] = {};
float w_sum = 0.f;
#pragma unroll
for (int n = 0; n < N_CHUNK; n++) {
if (n < an) {
w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E);
w_sum += w_n[n];
}
}
auto pass_B_body = [&](auto AN_TOK) {
constexpr int AN = decltype(AN_TOK)::value;
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
float2 corr2 = make_float2(corr, corr);
float2 a[2];
#pragma unroll
for (int j = 0; j < 2; j++) {
float2 old = make_float2(acc32[si * VEC + 2 * j],
acc32[si * VEC + 2 * j + 1]);
a[j] = float2_mul(old, corr2);
}
float2 f_cache[AN][2];
#pragma unroll
for (int n = 0; n < AN; n++) {
tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC,
f_cache[n]);
}
#pragma unroll
for (int n = 0; n < AN; n++) {
float2 wn = make_float2(w_n[n], w_n[n]);
#pragma unroll
for (int j = 0; j < 2; j++) {
a[j] = float2_fma(wn, f_cache[n][j], a[j]);
}
}
#pragma unroll
for (int j = 0; j < 2; j++) {
acc32[si * VEC + 2 * j] = a[j].x;
acc32[si * VEC + 2 * j + 1] = a[j].y;
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
float2 corr2 = make_float2(corr, corr);
float2 a[VEC / 2];
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
float2 old = make_float2(acc32[si * VEC + 2 * j],
acc32[si * VEC + 2 * j + 1]);
a[j] = float2_mul(old, corr2);
}
float2 f_cache[AN][VEC / 2];
#pragma unroll
for (int n = 0; n < AN; n++) {
tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]);
}
#pragma unroll
for (int n = 0; n < AN; n++) {
float2 wn = make_float2(w_n[n], w_n[n]);
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
a[j] = float2_fma(wn, f_cache[n][j], a[j]);
}
}
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
acc32[si * VEC + 2 * j] = a[j].x;
acc32[si * VEC + 2 * j + 1] = a[j].y;
}
}
};
if constexpr (NC == 4) {
switch (an) {
case 4:
pass_B_body(std::integral_constant<int, 4>{});
break;
case 3:
pass_B_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_B_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_B_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else if constexpr (NC == 3) {
switch (an) {
case 3:
pass_B_body(std::integral_constant<int, 3>{});
break;
case 2:
pass_B_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_B_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
} else {
static_assert(NC == 2);
switch (an) {
case 2:
pass_B_body(std::integral_constant<int, 2>{});
break;
case 1:
pass_B_body(std::integral_constant<int, 1>{});
break;
default:
__builtin_unreachable();
}
}
s_running = s_running * corr + w_sum;
m_running = m_new;
}
float inv_s = 1.f / s_running;
bf16_t* out_ptr = output + (long long)tb * H;
float2 output_sq_pair = {};
// When output RMSNorm is fused, the softmax denominator cancels:
// (acc / s) * rsqrt(mean((acc / s)^2) + eps)
// = acc * rsqrt(mean(acc^2) + eps * s^2).
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
uint2 packed;
auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
float2 inv2 = make_float2(inv_s, inv_s);
#pragma unroll
for (int j = 0; j < 2; j++) {
float2 old = make_float2(acc32[si * VEC + 2 * j],
acc32[si * VEC + 2 * j + 1]);
if constexpr (HAS_OUTPUT_NORM) {
output_sq_pair = float2_fma(old, old, output_sq_pair);
} else {
float2 mixed = float2_mul(old, inv2);
ov2[j] = __float22bfloat162_rn(mixed);
}
}
if constexpr (!HAS_OUTPUT_NORM) {
*reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
}
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
uint4 packed;
auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
float2 inv2 = make_float2(inv_s, inv_s);
#pragma unroll
for (int j = 0; j < VEC / 2; j++) {
float2 old =
make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]);
if constexpr (HAS_OUTPUT_NORM) {
output_sq_pair = float2_fma(old, old, output_sq_pair);
} else {
float2 mixed = float2_mul(old, inv2);
ov2[j] = __float22bfloat162_rn(mixed);
}
}
if constexpr (!HAS_OUTPUT_NORM) {
*reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
}
}
if constexpr (HAS_OUTPUT_NORM) {
if constexpr (OUTPUT_NORM_IN_SMEM) {
// The immutable weight copy is acquired once, at its first use.
if (tb == blockIdx.x) {
mbarrier_wait(plan.bar_output_norm_ready, 0);
}
}
float output_sq = output_sq_pair.x + output_sq_pair.y;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset);
}
if (lane == 0) {
plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f);
}
named_barrier_sync(CONSUMER_THREADS, 0);
float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f;
#pragma unroll
for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
total_sq +=
__shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS);
}
if (lane == 0) {
total_sq =
rsqrtf(total_sq / H + output_norm_eps * s_running * s_running);
}
float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0);
#pragma unroll
for (int si = 0; si < SLICES_PER_GROUP; si++) {
if constexpr (H == 7168) {
if (si == SLICES_PER_GROUP - 1) {
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
uint2 packed;
auto* values = reinterpret_cast<bf16_t*>(&packed);
#pragma unroll
for (int j = 0; j < 4; j++) {
const bf16_t* weight_ptr =
OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
float weight = __bfloat162float(weight_ptr[h_base + j]);
values[j] = __float2bfloat16(acc32[si * VEC + j] *
output_rsigma * weight);
}
*reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
continue;
}
}
int dt = si * CONSUMER_GROUPS + group;
if (dt >= NHT) continue;
int h_base = dt * K_TILE + k_local;
uint4 packed;
auto* values = reinterpret_cast<bf16_t*>(&packed);
#pragma unroll
for (int j = 0; j < VEC; j++) {
const bf16_t* weight_ptr =
OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
float weight = __bfloat162float(weight_ptr[h_base + j]);
values[j] =
__float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight);
}
*reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
}
}
}
}
cudaTriggerProgrammaticLaunchCompletion();
__syncthreads();
if (wid == 1) {
tmem_free(plan.tmem_base, TMEM_COLS_ALLOC);
}
#else
if (threadIdx.x == 0) {
printf("attn_res_fwd_online_v2_kernel requires sm_10x\n");
}
#endif
}
template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
bool OUTPUT_NORM_IN_SMEM = false>
static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
const bf16_t* delta, const bf16_t* res_weight,
const bf16_t* rms_weight, bf16_t* output, int N, int T,
int B, float rms_eps, int num_sm, cudaStream_t stream,
const bf16_t* output_norm_weight = nullptr,
float output_norm_eps = 0.f, int block_stride_m = 0,
int block_stride_r = 0) {
constexpr size_t smem_size =
((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) +
(OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) +
sizeof(FwdSmemPlan<NC>) + 15) &
~size_t(15);
auto kernel =
&attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA,
HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
static bool attrs_set = false;
if (!attrs_set) {
if (smem_size > 48 * 1024) {
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
}
attrs_set = true;
}
int grid = RELEASE_TMEM ? num_sm * 2 : num_sm;
cudaLaunchConfig_t config{};
config.gridDim = grid;
config.blockDim = BLK;
config.dynamicSmemBytes = smem_size;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.attrs = attrs;
config.numAttrs = 1;
cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
res_weight, rms_weight, output, N, T, B, block_stride_m,
block_stride_r, rms_eps, output_norm_weight,
output_norm_eps);
}
} // namespace fwd_prod_v2
} // namespace sm100
void kimi_k3_attn_res(torch::stable::Tensor& prefix,
torch::stable::Tensor const& delta,
torch::stable::Tensor const& blocks,
torch::stable::Tensor const& norm_weight,
torch::stable::Tensor const& qk_weight,
torch::stable::Tensor const& output_norm_weight,
torch::stable::Tensor& output, int64_t num_blocks,
double eps, double output_norm_eps) {
int const num_tokens = static_cast<int>(prefix.size(0));
int const device = prefix.get_device_index();
torch::stable::accelerator::DeviceGuard const device_guard(device);
cudaDeviceProp const* properties = get_device_prop();
STD_TORCH_CHECK(properties->major == 10,
"Kimi K3 AttnRes requires the SM100 family");
using namespace sm100::fwd_prod_v2;
// Two-source chunks and two resident CTAs are beneficial once setup is
// amortized by the long, full eight-block prefill workload.
if (num_blocks == 8 && num_tokens >= 4096) {
launch_fwd<7168, 2, true, true, true, true>(
static_cast<bf16_t const*>(blocks.data_ptr()),
static_cast<bf16_t*>(prefix.data_ptr()),
static_cast<bf16_t const*>(delta.data_ptr()),
static_cast<bf16_t const*>(qk_weight.data_ptr()),
static_cast<bf16_t const*>(norm_weight.data_ptr()),
static_cast<bf16_t*>(output.data_ptr()),
static_cast<int>(num_blocks) + 1, num_tokens, 1,
static_cast<float>(eps), properties->multiProcessorCount,
get_current_cuda_stream(device),
static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
static_cast<int>(blocks.stride(1)));
} else {
launch_fwd<7168, 4, false, true, true, true>(
static_cast<bf16_t const*>(blocks.data_ptr()),
static_cast<bf16_t*>(prefix.data_ptr()),
static_cast<bf16_t const*>(delta.data_ptr()),
static_cast<bf16_t const*>(qk_weight.data_ptr()),
static_cast<bf16_t const*>(norm_weight.data_ptr()),
static_cast<bf16_t*>(output.data_ptr()),
static_cast<int>(num_blocks) + 1, num_tokens, 1,
static_cast<float>(eps), properties->multiProcessorCount,
get_current_cuda_stream(device),
static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
static_cast<int>(blocks.stride(1)));
}
cudaError_t const error = cudaGetLastError();
STD_TORCH_CHECK(
error == cudaSuccess,
"Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error));
}
csrc/libtorch_stable/layernorm_kernels.cu
+10 -4
View File
@@ -251,7 +251,9 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
// For large num_tokens, use smaller blocks to increase SM concurrency.
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
dim3 grid(num_tokens);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -328,8 +330,13 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
/* This kernel is memory-latency bound in many scenarios.
When num_tokens is large, a smaller block size allows
for increased block occupancy on CUs and better latency
hiding on global mem ops. */
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
hiding on global mem ops. In batch-invariant mode the block size must
not depend on num_tokens, otherwise the same token would use a different
reduction width (and thus a different floating-point summation order)
across batches; lock it to 1024 to keep results bit-exact. */
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
dim3 block(std::min(hidden_size, max_block_size));
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -341,7 +348,6 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
bool offsets_are_multiple_of_vector_width =
hidden_size % vector_width == 0 && input_stride % vector_width == 0 &&
residual_stride % vector_width == 0;
bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const bool has_weight = weight.has_value();
if (has_weight) {
auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());
csrc/libtorch_stable/layernorm_quant_kernels.cu
+6 -3
View File
@@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant(
int num_tokens = input.numel() / hidden_size;
// For large num_tokens, use smaller blocks to increase SM concurrency.
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
dim3 grid(num_tokens);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant(
When num_tokens is large, a smaller block size allows
for increased block occupancy on CUs and better latency
hiding on global mem ops. */
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
dim3 block(std::min(hidden_size, max_block_size));
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant(
auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
bool ptrs_are_aligned =
inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&
!batch_invariant_launch) {
LAUNCH_FUSED_ADD_RMS_NORM(8);
csrc/libtorch_stable/ops.h
+11
View File
@@ -315,6 +315,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
std::optional<torch::stable::Tensor> index_q_out,
const std::string& kv_cache_dtype, bool skip_index_branch);
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
void kimi_k3_attn_res(torch::stable::Tensor& prefix,
torch::stable::Tensor const& delta,
torch::stable::Tensor const& blocks,
torch::stable::Tensor const& norm_weight,
torch::stable::Tensor const& qk_weight,
torch::stable::Tensor const& output_norm_weight,
torch::stable::Tensor& output, int64_t num_blocks,
double eps, double output_norm_eps);
#endif
// Sampler kernels (shared CUDA/ROCm)
void apply_repetition_penalties_(
torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+10 -42
View File
@@ -49,11 +49,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
cudaGridDependencySynchronize();
cudaTriggerProgrammaticLaunchCompletion();
#endif
// Precompute SF layout parameter (constant for entire kernel).
int32_t const numKTiles = (outputCols + 63) / 64;
@@ -128,11 +123,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
cudaGridDependencySynchronize();
cudaTriggerProgrammaticLaunchCompletion();
#endif
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
@@ -211,8 +201,6 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
auto stream = get_current_cuda_stream(input.get_device_index());
auto* device_props = get_device_prop();
int const sm_version = device_props->major * 10 + device_props->minor;
int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE);
@@ -236,21 +224,10 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
input.scalar_type(), "nvfp4_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
cudaLaunchConfig_t config = {};
config.gridDim = grid;
config.blockDim = block;
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
config.attrs = attrs;
cudaLaunchKernelEx(&config, vllm::cvt_fp16_to_fp4<cuda_type, false>,
m, n, output_n, num_padded_cols, input_ptr,
input_sf_ptr,
reinterpret_cast<uint32_t*>(output_ptr),
reinterpret_cast<uint32_t*>(sf_out));
vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
m, n, output_n, num_padded_cols, input_ptr, input_sf_ptr,
reinterpret_cast<uint32_t*>(output_ptr),
reinterpret_cast<uint32_t*>(sf_out));
});
} else {
int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;
@@ -263,21 +240,12 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
input.scalar_type(), "nvfp4_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
cudaLaunchConfig_t config = {};
config.gridDim = grid;
config.blockDim = block;
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
config.attrs = attrs;
cudaLaunchKernelEx(
&config, vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>, m, n,
output_n, output_sf_n_unpadded, num_packed_cols, input_ptr,
input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
reinterpret_cast<uint32_t*>(sf_out));
vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
<<<grid, block, 0, stream>>>(
m, n, output_n, output_sf_n_unpadded, num_packed_cols,
input_ptr, input_sf_ptr,
reinterpret_cast<uint32_t*>(output_ptr),
reinterpret_cast<uint32_t*>(sf_out));
});
}
}
csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+4 -1
View File
@@ -2,6 +2,7 @@
#include "../../torch_utils.h"
#include "../../dispatch_utils.h"
#include "../../../core/batch_invariant.hpp"
#include "layernorm_utils.cuh"
#include "quant_conversions.cuh"
@@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch(
auto num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
const int max_block_size = (num_tokens <= 256) ? 512 : 256;
const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
const int max_block_size =
batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256);
dim3 block(std::min(hidden_size, max_block_size));
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
csrc/libtorch_stable/torch_bindings.cpp
+11 -4
View File
@@ -330,10 +330,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
// conditionally compiled so impl registration is in source file
ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
// BF16 skinny GEMM (M<=32, weight-BW-bound shapes, e.g. MTP eh_proj).
// conditionally compiled so impl registration is in source file
ops.def("bf16_skinny_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, "
@@ -472,6 +468,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"int block_size, Tensor!? q_out, Tensor!? index_q_out, "
"str kv_cache_dtype, bool skip_index_branch=False) -> ()");
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
ops.def(
"kimi_k3_attn_res("
"Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
"Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
"int num_blocks, float eps, float output_norm_eps) -> ()");
#endif
// Apply repetition penalties to logits in-place.
ops.def(
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -697,6 +701,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
#endif
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
#endif
// Sampler kernels (shared CUDA/ROCm)
ops.impl("apply_repetition_penalties_",
csrc/quickreduce/quick_reduce.h
+30 -8
View File
@@ -22,17 +22,25 @@ template <typename AllReduceKernel, typename T>
__global__ __quickreduce_launch_bounds_two_shot__ static void
allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
int rank, uint8_t** dbuffer_list,
uint32_t data_offset, uint32_t flag_color,
uint32_t data_offset, uint32_t* d_flag_counters,
int64_t data_size_per_phase) {
int block = blockIdx.x;
int grid = gridDim.x;
// Load this block's counter from device memory and advance it on-device,
// so the color keeps changing across graph replays instead of being frozen.
uint32_t flag_color = d_flag_counters[blockIdx.x];
while (block < num_blocks) {
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
flag_color, data_size_per_phase);
block += grid;
flag_color++;
}
// All threads compute the same final value; one writer per block is enough.
if (threadIdx.x == 0 && threadIdx.y == 0) {
d_flag_counters[blockIdx.x] = flag_color;
}
}
#define TWOSHOT_DISPATCH(__codec) \
@@ -42,21 +50,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
flag_color, this->kMaxProblemSize); \
d_flag_counters, this->kMaxProblemSize); \
} else if (world_size == 4) { \
using LineCodec = __codec<T, 4>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
flag_color, this->kMaxProblemSize); \
d_flag_counters, this->kMaxProblemSize); \
} else if (world_size == 8) { \
using LineCodec = __codec<T, 8>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
flag_color, this->kMaxProblemSize); \
d_flag_counters, this->kMaxProblemSize); \
}
// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
@@ -69,7 +77,7 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
flag_color, this->kMaxProblemSize); \
d_flag_counters, this->kMaxProblemSize); \
} else { \
throw std::runtime_error( \
"INT3 quick all-reduce is only supported for world_size == 2 " \
@@ -94,7 +102,7 @@ struct DeviceComms {
static int constexpr kMaxWorldSize = 8;
bool initialized = false;
uint32_t flag_color = 1;
uint32_t* d_flag_counters = nullptr;
int world_size;
int rank;
@@ -128,6 +136,16 @@ struct DeviceComms {
// Clear the flags buffer.
HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
// One flag-color counter per block, advanced by the kernel. Start at 1
// to stay clear of the flags buffer we just zeroed.
HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
{
std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
kMaxNumBlocks * sizeof(uint32_t),
hipMemcpyHostToDevice));
}
// Device-side list of IPC buffers.
buffer_list.resize(world_size);
HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
@@ -144,6 +162,12 @@ struct DeviceComms {
hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
void destroy() {
// Allocated before `initialized` flips true, so free it on its own guard
// to avoid a leak if init fails partway through.
if (d_flag_counters) {
HIP_CHECK(hipFree(d_flag_counters));
d_flag_counters = nullptr;
}
if (initialized) {
for (int i = 0; i < world_size; i++) {
if (i != rank) {
@@ -211,8 +235,6 @@ struct DeviceComms {
break;
}
HIP_CHECK(cudaGetLastError());
// Rotate the flag color.
flag_color += divceil(N, grid);
}
};
docker/Dockerfile
+25 -8
View File
@@ -25,6 +25,10 @@
ARG CUDA_VERSION=13.0.3
ARG PYTHON_VERSION=3.12
ARG UBUNTU_VERSION=22.04
# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
# This version is only used for CUDA 13+ builds; CUDA 12 falls back to
# the default NCCL version shipped with the base image.
ARG NCCL_VERSION=2.30.7
# By parameterizing the base images, we allow third-party to use their own
# base images. One use case is hermetic builds with base images stored in
@@ -477,10 +481,17 @@ WORKDIR /workspace
# Build DeepEP wheels
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
# Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
ARG DEEPEP_COMMIT_HASH=73b6ea4
ARG DEEPEP_COMMIT_HASH=d4f41e4e93
ARG NVSHMEM_VER
ARG NCCL_VERSION
RUN --mount=type=cache,target=/opt/uv/cache \
mkdir -p /tmp/ep_kernels_workspace/dist && \
CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
> /tmp/nccl-override.txt && \
export UV_OVERRIDE=/tmp/nccl-override.txt; \
fi && \
export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
/tmp/install_python_libraries.sh \
--workspace /tmp/ep_kernels_workspace \
@@ -644,6 +655,7 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
ARG NCCL_VERSION
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL
@@ -696,7 +708,6 @@ RUN apt-get update -y \
# Install CUDA development tools for runtime JIT compilation
# (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \
apt-get update -y && \
apt-get install -y --no-install-recommends --allow-change-held-packages \
cuda-nvcc-${CUDA_VERSION_DASH} \
@@ -709,12 +720,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
libnuma-dev \
# numactl CLI for NUMA binding at runtime
numactl && \
# Fixes nccl_allocator requiring nccl.h at runtime
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
# so we pin the version to match our CUDA version
NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \
apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
rm -rf /var/lib/apt/lists/*
# Install uv for faster pip installs
@@ -734,6 +739,18 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
&& chgrp -R 0 /opt/uv \
&& chmod -R g+rwX,a+rX /opt/uv
# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins
# an older version as a transitive dep; this override forces uv to use our
# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op).
RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
if [ "$CUDA_MAJOR" -ge 13 ]; then \
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
> /etc/uv-overrides.txt; \
else \
touch /etc/uv-overrides.txt; \
fi
ENV UV_OVERRIDE=/etc/uv-overrides.txt
# ----------------------------------------------------------------------
# Non-root support (opt-in)
# ----------------------------------------------------------------------
docker/Dockerfile.rocm_base
+1 -1
View File
@@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
ENV AITER_ROCM_ARCH=gfx942;gfx950
ENV MORI_GPU_ARCHS=gfx942;gfx950
# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime
# Required for RCCL in ROCm7.1
ENV HSA_NO_SCRATCH_RECLAIM=1
docker/versions.json
+4 -1
View File
@@ -10,6 +10,9 @@
"UBUNTU_VERSION": {
"default": "22.04"
},
"NCCL_VERSION": {
"default": "2.30.7"
},
"BUILD_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
},
@@ -56,7 +59,7 @@
"default": "cuda"
},
"DEEPEP_COMMIT_HASH": {
"default": "73b6ea4"
"default": "d4f41e4e93"
},
"GIT_REPO_CHECK": {
"default": "0"
docs/.nav.yml
+3 -1
View File
@@ -56,7 +56,9 @@ nav:
- API Reference:
- api/README.md
- api/vllm
- CLI Reference: cli
- CLI Reference:
- cli/README.md
- vllm: cli
- Community:
- community/*
- Governance: governance
docs/cli/.nav.yml
+7 -9
View File
@@ -1,10 +1,8 @@
nav:
- README.md
- serve.md
- chat.md
- complete.md
- run-batch.md
- vllm bench:
- bench/**/*.md
- vllm launch:
- launch/**/*.md
- "*.md"
- bench:
- bench/*.md
- sweep:
- bench/sweep/*.md
- launch:
- launch/*.md
docs/cli/bench/latency.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench latency
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_latency.inc.md"
docs/cli/bench/mm_processor.md
-55
View File
@@ -1,55 +0,0 @@
# vllm bench mm-processor
## Overview
`vllm bench mm-processor` profiles the multimodal input processor pipeline of
vision-language models. It measures per-stage latency from the HuggingFace
processor through to the encoder forward pass, helping you identify
preprocessing bottlenecks and understand how different image resolutions or
item counts affect end-to-end request time.
The benchmark supports two data sources: synthetic random multimodal inputs
(`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before
measurement to ensure stable results.
## Quick Start
```bash
vllm bench mm-processor \
--model Qwen/Qwen2-VL-7B-Instruct \
--dataset-name random-mm \
--num-prompts 50 \
--random-input-len 300 \
--random-output-len 40 \
--random-mm-base-items-per-request 2 \
--random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
--random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}'
```
## Measured Stages
| Stage | Description |
| ----- | ----------- |
| `get_mm_hashes_secs` | Time spent hashing multimodal inputs |
| `get_cache_missing_items_secs` | Time spent looking up the processor cache |
| `apply_hf_processor_secs` | Time spent in the HuggingFace processor |
| `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs |
| `apply_prompt_updates_secs` | Time spent updating prompt tokens |
| `preprocessor_total_secs` | Total preprocessing time |
| `encoder_forward_secs` | Time spent in the encoder model forward pass |
| `num_encoder_calls` | Number of encoder invocations per request |
The benchmark also reports end-to-end latency (TTFT + decode time) per
request. Use `--metric-percentiles` to select which percentiles to report
(default: p99) and `--output-json` to save results.
For more examples (HF datasets, warmup, JSON output), see
[Benchmarking CLI — Multimodal Processor Benchmark](../../benchmarking/cli.md#multimodal-processor-benchmark).
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_mm_processor.inc.md"
docs/cli/bench/serve.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_serve.inc.md"
docs/cli/bench/sweep/plot.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench sweep plot
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_sweep_plot.inc.md"
docs/cli/bench/sweep/plot_pareto.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench sweep plot_pareto
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_sweep_plot_pareto.inc.md"
docs/cli/bench/sweep/serve.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench sweep serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_sweep_serve.inc.md"
docs/cli/bench/sweep/serve_workload.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench sweep serve_workload
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_sweep_serve_workload.inc.md"
docs/cli/bench/throughput.md
-9
View File
@@ -1,9 +0,0 @@
# vllm bench throughput
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/bench_throughput.inc.md"
docs/cli/chat.md
-5
View File
@@ -1,5 +0,0 @@
# vllm chat
## Arguments
--8<-- "docs/generated/argparse/chat.inc.md"
docs/cli/complete.md
-5
View File
@@ -1,5 +0,0 @@
# vllm complete
## Arguments
--8<-- "docs/generated/argparse/complete.inc.md"
docs/cli/json_tip.inc.md
-10
View File
@@ -1,10 +0,0 @@
<!-- markdownlint-disable MD041 -->
When passing JSON CLI arguments, the following sets of arguments are equivalent:
- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'`
- `--json-arg.key1 value1 --json-arg.key2.key3 value2`
Additionally, list elements can be passed individually using `+`:
- `--json-arg '{"key4": ["value3", "value4", "value5"]}'`
- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'`
docs/cli/launch/render.md
-22
View File
@@ -1,22 +0,0 @@
# vllm launch render
## Overview
`vllm launch render` starts a GPU-less rendering server for preprocessing and
postprocessing only.
```bash
vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100
```
This command reuses the standard serving parser, so model, frontend,
networking, and related CLI options follow the same conventions as
[`vllm serve`](../serve.md).
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/launch_render.inc.md"
docs/cli/run-batch.md
-9
View File
@@ -1,9 +0,0 @@
# vllm run-batch
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/run-batch.inc.md"
docs/cli/serve.md
-9
View File
@@ -1,9 +0,0 @@
# vllm serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse/serve.inc.md"
docs/configuration/engine_args.md
+1 -9
View File
@@ -11,12 +11,4 @@ Engine arguments control the behavior of the vLLM engine.
The engine argument classes, [EngineArgs][vllm.engine.arg_utils.EngineArgs] and [AsyncEngineArgs][vllm.engine.arg_utils.AsyncEngineArgs], are a combination of the configuration classes defined in [vllm.config][]. Therefore, if you are interested in developer documentation, we recommend looking at these configuration classes as they are the source of truth for types, defaults and docstrings.
--8<-- "docs/cli/json_tip.inc.md"
## `EngineArgs`
--8<-- "docs/generated/argparse/engine_args.inc.md"
## `AsyncEngineArgs`
--8<-- "docs/generated/argparse/async_engine_args.inc.md"
--8<-- "gen:engine-args"
docs/contributing/model/transcription.md
+1 -1
View File
@@ -195,7 +195,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
The API server takes care of basic audio I/O and optional chunking before building prompts:
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into chunks and generates a prompt per chunk. There is no overlap between chunks, overlap_chunk_second controls the size of the search window used to find the split point.
- Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
Relevant server logic:
docs/deployment/docker.md
+20
View File
@@ -8,6 +8,26 @@ toc_depth: 2
--8<-- "docs/getting_started/installation/gpu.md:pre-built-images"
## Persist the compile cache across containers
Mounting the Hugging Face cache keeps model weights across containers, but each
new container still starts with an empty `VLLM_CACHE_ROOT` (default
`~/.cache/vllm`) and recompiles the model's `torch.compile` artifacts. Mount a
named volume at that path to reuse the inductor, Triton, and AOT artifacts from
the second container onward:
```bash
docker run --rm --gpus all \
-v ~/.cache/huggingface:/root/.cache/huggingface \
-v vllm-cache:/root/.cache/vllm \
-p 8000:8000 \
vllm/vllm-openai:latest \
meta-llama/Llama-3.1-8B-Instruct
```
See [Faster Startup](../configuration/optimization.md#faster-startup) for the
mechanism and for what invalidates the cache.
## Run as a non-root user
The CUDA `vllm/vllm-openai` image runs as root by default for backward
docs/deployment/integrations/llm-d.md
+34 -2
View File
@@ -1,5 +1,37 @@
# llm-d
vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers.
[llm-d](https://llm-d.ai/) is a Kubernetes-native distributed inference framework for serving large language models at scale, with vLLM as its primary inference engine. llm-d coordinates a fleet of vLLM instances across a cluster so that performance holds up under real production traffic, achieving the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators.
You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d-to-the-cncf-evolving-kubernetes-into-sota-ai-infrastructure/) founded by Red Hat, Google Cloud, IBM Research, CoreWeave, and NVIDIA.
## What llm-d adds to vLLM
A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own:
- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it.
- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM.
- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts.
- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput.
- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch.
These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear.
## Performance
Representative benchmarked results across accelerators:
- **3x higher output throughput** and **2x faster TTFT** from prefix-aware routing vs round-robin (Llama 3.1 70B, AMD MI300X)
- **Up to 70% higher tokens/sec** from prefill/decode disaggregation (GPT-OSS, NVIDIA B200)
- **13.9x throughput** from hierarchical KV offloading at high concurrency vs GPU-only (NVIDIA H100)
See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and reproducible benchmarks on [Prism](https://prism.llm-d.ai/).
## Get started
1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration.
2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload.
3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment.
You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
Questions and contributions are welcome on [GitHub](https://github.com/llm-d/llm-d) and [Slack](https://llm-d.ai/slack).
docs/design/attention_backends.md
+14 -90
View File
@@ -1,15 +1,9 @@
# Attention Backend Feature Support
This document is auto-generated by `tools/pre_commit/generate_attention_backend_docs.py`.
It shows the feature support for each registered attention backend
based on the checks in `AttentionBackend.validate_configuration()`.
**Do not edit this file manually.** Run the following command to
regenerate it:
```bash
python tools/pre_commit/generate_attention_backend_docs.py
```
The priority and feature tables on this page are auto-generated from the
attention backend registry by
`docs/mkdocs/gen_files/generate_attention_backends.py`, based on the checks in
`AttentionBackend.validate_configuration()`.
## Setting the Attention Backend
@@ -98,40 +92,11 @@ Priority is **1 = highest** (tried first).
### Standard Attention (MHA, MQA, GQA)
**Blackwell (SM 10.x):**
| Priority | Backend |
| -------- | ------- |
| 1 | `FLASHINFER` |
| 2 | `FLASH_ATTN` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
**Ampere/Hopper (SM 8.x-9.x):**
| Priority | Backend |
| -------- | ------- |
| 1 | `FLASH_ATTN` |
| 2 | `FLASHINFER` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
--8<-- "gen:priority-standard"
### MLA Attention (DeepSeek-style)
**Blackwell (SM 10.x):**
| Priority | Backend |
| -------- | ------- |
| 1 | `FLASHINFER_MLA` |
| 2 | `TOKENSPEED_MLA` |
| 3 | `CUTLASS_MLA` |
| 4 | `FLASH_ATTN_MLA` |
| 5 | `FLASHMLA` |
| 6 | `TRITON_MLA` |
| 7 | `FLASHINFER_MLA_SPARSE`**\*** |
| 8 | `FLASHMLA_SPARSE` |
--8<-- "gen:priority-mla"
> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
>
@@ -157,24 +122,7 @@ Priority is **1 = highest** (tried first).
## Standard Attention (MHA, MQA, GQA) Backends
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
--8<-- "gen:table-standard"
> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
>
@@ -188,9 +136,7 @@ automatic priority lists above. A lightning indexer scores KV blocks, the
top-k blocks (plus fixed init/local blocks) are selected, and attention
attends only to those blocks; index keys live in a separate side cache.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
--8<-- "gen:table-minimax"
## MLA (Multi-head Latent Attention) Backends
@@ -203,38 +149,20 @@ To explicitly select a prefill backend, use
Otherwise, the prefill backend is selected automatically at runtime based on
hardware and configuration.
| Backend | Description | Dtypes | Compute Cap. | Notes |
| ------- | ----------- | ------ | ------------ | ----- |
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
--8<-- "gen:table-mla-prefill"
> **‡** Automatic selection tries FlashAttention first. On Blackwell
> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64,
> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other
> GPUs, only FlashAttention is considered.
### Decode Backends
MLA decode backends are selected using the standard
`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
--8<-- "gen:table-mla-decode"
### DeepSeek V4 Decode Backends
@@ -245,8 +173,4 @@ pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and
`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x |
| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
--8<-- "gen:table-mla-v4-decode"
docs/design/cuda_graphs_multimodal.md
+2 -1
View File
@@ -101,7 +101,7 @@ When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks
Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
!!! note
Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
Video CUDA graphs are automatically disabled when video token pruning (EVS or VidCom2) is enabled, since pruning makes the token count data-dependent and incompatible with CUDA graph capture.
Mixed inputs (image+video) per prompt are also supported now.
@@ -129,6 +129,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ |
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
docs/design/custom_op.md
-2
View File
@@ -122,8 +122,6 @@ For example:
--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"
--8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer"
--8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
```
docs/features/kv_offloading_usage.md
+66 -1
View File
@@ -157,7 +157,7 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti
The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required.
PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes.
The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected.
| Key | Required | Default | Notes |
| --- | --- | --- | --- |
@@ -176,6 +176,71 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
- `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide.
- `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision.
#### Orchestration-Layer Protocol
The P2P tier does not decide *which* peer to pull from — that is the orchestration layer's job (the router/EPP and its scheduler). The orchestrator drives every transfer through a request's `kv_transfer_params` dict: it picks the request's role, allocates a unique transaction ID, and supplies the remote peer's address. All block lookup, hash matching, and NIXL transfer happen at the tier level below; the orchestrator only sets the correct role keys and enforces the allowed combinations.
Every vLLM instance is a symmetric **peer**. Per request it acts as a **consumer** (pulls KV blocks from a remote peer's CPU cache instead of computing locally) or a **producer** (serves blocks from its own CPU cache to remote consumers) — or both, on the same session, for different requests. Roles are chosen per request by the keys below; there are no fixed prefiller/decoder processes.
Three role keys are defined, each mapping to a sub-dict. All are optional; a request with none of them uses the tier only as a local CPU cache.
Each key names the **remote counterpart** this peer transfers with (not this
peer's own role), so the name reads as "the remote ___ I transfer with".
| Key | Set on | Value fields | Meaning |
| --- | --- | --- | --- |
| `remote_decoder` | prefill producer request | `kv_request_id` | Peer computes KV and keeps it available in CPU cache for the remote decoder to pull. |
| `remote_prefiller` | decode consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer pulls KV from the remote prefiller at the given address (classic P/D disaggregation). |
| `remote_kv_source` | P2P consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer looks up and pulls whatever blocks the remote source currently holds in CPU cache. |
Field semantics:
- `kv_request_id` (str): unique transaction ID allocated by the orchestrator and pushed to every peer involved in the transfer; used to correlate the lookup, fetch, and transfer-done messages. The producer is implicit — it serves whatever block hashes it currently holds in its CPU cache for that ID.
- `remote_host` (str): IP/hostname of the remote peer's control socket to query. Must be the peer's routable node IP (see [Environment Variables](#environment-variables)).
- `remote_port` (int): the peer's bound control-socket port, i.e. `base + data_parallel_index` for the selected DP rank.
Allowed and forbidden combinations:
- **`remote_decoder` + `remote_kv_source`** is the only legal multi-key combination: a prefill producer may *also* act as a P2P consumer for the same request — skipping prefix prefill by pulling cached blocks from a source while still keeping its own computed blocks available for a downstream decoder.
- Forbidden: `remote_prefiller` + `remote_decoder` (contradictory roles), `remote_prefiller` + `remote_kv_source` (two competing fetch sources), and all three together.
Minimal examples (values that would appear in the request's `kv_transfer_params`):
```python
# Prefill producer — compute and keep KV for a remote decoder to pull
kv_transfer_params = {"remote_decoder": {"kv_request_id": "<unique-transfer-id>"}}
# Decode consumer — pull KV from a specific prefiller (classic P/D)
kv_transfer_params = {
"remote_prefiller": {
"kv_request_id": "<unique-transfer-id>",
"remote_host": "<prefiller-node-ip>",
"remote_port": 5710,
}
}
# P2P consumer — pull whatever the source already has cached
kv_transfer_params = {
"remote_kv_source": {
"kv_request_id": "<unique-transfer-id>",
"remote_host": "<source-node-ip>",
"remote_port": 5710,
}
}
```
Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the keys above:
1. Both peers already have listener threads on their control sockets (see [Environment Variables](#environment-variables)).
2. **Lookup.** The consumer's tiering manager does per-block lookups; in P2P mode the tier returns `None` and registers the key. At `on_schedule_end` the consumer sends one **`LookupMsg`** (`kv_request_id` + block hashes) to the peer, per request, per step.
3. The producer matches those hashes against its local CPU cache and replies with a **`LookupRespMsg`** carrying the hit block hashes.
4. **Resolve.** Retried lookups now return hit / miss / in-flight. The consumer calls `submit_load` for hits only, allocating CPU slots only for hits.
5. The consumer sends a **`FetchMsg`** (`kv_request_id`, block hashes, destination block indexes).
6. The producer performs the **NIXL WRITE** transfer and sends **`TransferDone`** with a success status.
7. On `get_finished`, hits are loaded into GPU as ordinary cache hits; misses are recomputed by the engine.
In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached.
## Tuning Tips
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
docs/features/multimodal_inputs.md
+46 -4
View File
@@ -350,6 +350,31 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
#### Video Token Pruning
For supported models, vLLM can prune video tokens after the vision encoder to
reduce prefill time and KV cache usage, at some cost in accuracy. Set
`--video-pruning-rate <q>` to prune the fraction `q` of video tokens from each
video, and `--video-pruning-method` to choose the training-free algorithm:
- **`evs`** (Efficient Video Sampling, default): drops the tokens with the
lowest temporal dissimilarity to the previous frame. The first frame is
always fully retained.
- **`vidcom2`** (Video Compression Commander): scores tokens by similarity to
video-level and frame-level feature centers and gives distinctive frames a
larger share of the budget. At least one token per frame is retained.
```bash
vllm serve Qwen/Qwen3-VL-8B-Instruct \
--video-pruning-rate 0.75 --video-pruning-method vidcom2
```
!!! note
`evs` is supported by all models implementing multimodal pruning;
`vidcom2` is currently supported by Qwen3-VL only. Unsupported combinations
are rejected at startup. Enabling video pruning also disables encoder CUDA
graphs, since the retained token count becomes data-dependent.
### Audio Inputs
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
@@ -818,16 +843,18 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu
#### Video Decoding Backend
vLLM decodes video bytes into frames using a selectable decoding backend. Three
vLLM decodes video bytes into frames using a selectable decoding backend. Five
backends are supported:
- `opencv` (default): OpenCV-based decoder.
- `pyav`: PyAV decoder.
- `torchcodec`: TorchCodec (PyTorch-native) decoder.
- `pynvvideocodec`: NVIDIA NVDEC-based decoder.
- `deepstream`: NVIDIA DeepStream NVDEC-based decoder.
All three backends are ultimately backed by FFmpeg. `torchcodec` lets
you choose which FFmpeg version is used while `opencv` and `pyav` rely on
whichever FFmpeg build they were linked against.
The CPU backends are backed by FFmpeg. `torchcodec` lets you choose which FFmpeg
version is used while `opencv` and `pyav` rely on whichever FFmpeg build they
were linked against.
Select the backend by passing the `backend` parameter via `--media-io-kwargs`:
@@ -854,6 +881,21 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
--media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}'
```
**PyNvVideoCodec-specific parameters:**
- `hw_decoders`: Maximum number of concurrent hardware decoder slots retained
by each API server process. It must be a positive integer and defaults to `2`,
which is the recommended starting point for concurrent video workloads.
Because vLLM reserves GPU memory for these slots at startup, this value cannot
be overridden per request. Benchmark before increasing it because each
additional slot increases the GPU memory reservation.
```bash
# Example: explicitly use the recommended 2 hardware decoders
vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
--media-io-kwargs '{"video": {"backend": "pynvvideocodec", "hw_decoders": 2}}'
```
#### Video Frame Recovery
For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place.
docs/features/quantization/modelopt.md
+14
View File
@@ -19,6 +19,20 @@ following `quantization.quant_algo` values:
- `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`).
- `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`).
!!! note
For NVFP4 checkpoints, vLLM selects a GEMM kernel automatically at load
time from the backends available on the current platform (CUTLASS,
FlashInfer, Marlin, and others). On GPUs without a supported native FP4
GEMM kernel, vLLM falls back to weight-only (W4A16) execution via Marlin
and logs a warning; this may reduce throughput for compute-heavy
workloads. Use `--linear-backend` to override the automatic selection
(this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment
variable). Values relevant to NVFP4 include `cutlass`,
`flashinfer_cutlass`, `flashinfer_trtllm`, `flashinfer_cudnn`, and
`marlin`; the full list is documented under `KernelConfig` on the
[Engine Arguments](../../configuration/engine_args.md) page and shown by
`vllm serve --help=KernelConfig`.
## Quantizing HuggingFace Models with PTQ
You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory.
docs/mkdocs/hooks/generate_argparse.py → docs/mkdocs/gen_files/generate_argparse.py
+183 -48
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib.metadata
import importlib.util
import inspect
import logging
import sys
import textwrap
@@ -10,17 +11,21 @@ from argparse import SUPPRESS, Action, HelpFormatter
from collections.abc import Callable, Iterable
from importlib.machinery import ModuleSpec
from pathlib import Path
from typing import TYPE_CHECKING, Literal
from typing import TYPE_CHECKING
from unittest.mock import MagicMock, patch
import mkdocs_gen_files
import regex as re
from pydantic_core import core_schema
logger = logging.getLogger("mkdocs")
ROOT_DIR = Path(__file__).parent.parent.parent.parent
ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse"
sys.path.insert(0, str(ROOT_DIR))
sys.path.insert(0, str(Path(__file__).parent))
from generated_content import fill_markers # noqa: E402
def mock_if_no_torch(mock_module: str, mock: MagicMock):
@@ -132,8 +137,8 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100):
bench_latency = auto_mock("vllm.benchmarks", "latency")
bench_mm_processor = auto_mock("vllm.benchmarks", "mm_processor")
bench_serve = auto_mock("vllm.benchmarks", "serve")
bench_startup = auto_mock("vllm.benchmarks", "startup")
bench_sweep_plot = auto_mock("vllm.benchmarks.sweep.plot", "SweepPlotArgs")
bench_sweep_plot_pareto = auto_mock(
"vllm.benchmarks.sweep.plot_pareto", "SweepPlotParetoArgs"
@@ -142,12 +147,28 @@ bench_sweep_serve = auto_mock("vllm.benchmarks.sweep.serve", "SweepServeArgs")
bench_sweep_serve_workload = auto_mock(
"vllm.benchmarks.sweep.serve_workload", "SweepServeWorkloadArgs"
)
bench_sweep_startup = auto_mock("vllm.benchmarks.sweep.startup", "SweepStartupArgs")
bench_throughput = auto_mock("vllm.benchmarks", "throughput")
AsyncEngineArgs = auto_mock("vllm.engine.arg_utils", "AsyncEngineArgs")
EngineArgs = auto_mock("vllm.engine.arg_utils", "EngineArgs")
ChatCommand = auto_mock("vllm.entrypoints.cli.openai", "ChatCommand")
CompleteCommand = auto_mock("vllm.entrypoints.cli.openai", "CompleteCommand")
BenchmarkSubcommand = auto_mock(
"vllm.entrypoints.cli.benchmark.main", "BenchmarkSubcommand"
)
import_bench_subcommands = auto_mock(
"vllm.entrypoints.cli.benchmark.main", "_import_bench_subcommand_modules"
)
BenchmarkSubcommandBase = auto_mock(
"vllm.entrypoints.cli.benchmark.base", "BenchmarkSubcommandBase"
)
BenchmarkMMProcessorSubcommand = auto_mock(
"vllm.entrypoints.cli.benchmark.mm_processor", "BenchmarkMMProcessorSubcommand"
)
LaunchSubcommandBase = auto_mock("vllm.entrypoints.cli.launch", "LaunchSubcommandBase")
launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION")
RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand")
sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS")
openai_cli_args = auto_mock("vllm.entrypoints.openai", "cli_args")
openai_run_batch = auto_mock("vllm.entrypoints.openai", "run_batch")
@@ -179,7 +200,7 @@ class MarkdownFormatter(HelpFormatter):
def add_text(self, text: str):
if text:
self._markdown_output.append(f"{text.strip()}\n\n")
self._markdown_output.append(f"{inspect.cleandoc(text)}\n\n")
def add_usage(self, usage, actions, groups, prefix=None):
pass
@@ -241,49 +262,163 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser:
return _parser or parser
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
logger.info("Generating argparse documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
# Create the ARGPARSE_DOC_DIR if it doesn't exist
if not ARGPARSE_DOC_DIR.exists():
ARGPARSE_DOC_DIR.mkdir(parents=True)
# Create parsers to document
parsers = {
# Engine args
"engine_args": create_parser(EngineArgs.add_cli_args),
"async_engine_args": create_parser(
AsyncEngineArgs.add_cli_args, async_args_only=True
),
# CLI
"serve": create_parser(openai_cli_args.make_arg_parser),
"chat": create_parser(ChatCommand.add_cli_args),
"complete": create_parser(CompleteCommand.add_cli_args),
"launch_render": create_parser(RenderSubcommand.add_cli_args),
"run-batch": create_parser(openai_run_batch.make_arg_parser),
# Benchmark CLI
"bench_latency": create_parser(bench_latency.add_cli_args),
"bench_mm_processor": create_parser(bench_mm_processor.add_cli_args),
"bench_serve": create_parser(bench_serve.add_cli_args),
"bench_sweep_plot": create_parser(bench_sweep_plot.add_cli_args),
"bench_sweep_plot_pareto": create_parser(bench_sweep_plot_pareto.add_cli_args),
"bench_sweep_serve": create_parser(bench_sweep_serve.add_cli_args),
"bench_sweep_serve_workload": create_parser(
bench_sweep_serve_workload.add_cli_args
),
"bench_throughput": create_parser(bench_throughput.add_cli_args),
}
# Generate documentation for each parser
for stem, parser in parsers.items():
doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
# Specify encoding for building on Windows
with open(doc_path, "w", encoding="utf-8") as f:
f.write(super(type(parser), parser).format_help())
logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
def format_help(parser: FlexibleArgumentParser) -> str:
"""Format a parser's help as markdown using `MarkdownFormatter`."""
return super(type(parser), parser).format_help()
if __name__ == "__main__":
on_startup("build", False)
# Absolute docs URLs are kept in the help text because they are useful in the
# terminal. Wrap them as markdown links so the `url_schemes` hook can rewrite
# them into doc-relative links / cross-references at render time.
_DOCS_URL = re.compile(r"https://docs\.vllm\.ai/en/[^/\s]+/[^\s)>]+")
def linkify_docs_urls(text: str) -> str:
"""Wrap bare docs.vllm.ai URLs in help text as markdown links."""
return _DOCS_URL.sub(lambda m: f"[{m.group()}]({m.group()})", text)
logger.info("Generating argparse documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
# The JSON tip is always rendered immediately before generated argument content,
# and the generator is its only consumer, so it lives here rather than in a
# separate snippet file. (The runtime terminal equivalent is
# `FlexibleArgumentParser._json_tip` in vllm/utils/argparse_utils.py.)
JSON_TIP = """## JSON CLI Arguments
When passing JSON CLI arguments, the following sets of arguments are equivalent:
- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'`
- `--json-arg.key1 value1 --json-arg.key2.key3 value2`
Additionally, list elements can be passed individually using `+`:
- `--json-arg '{"key4": ["value3", "value4", "value5"]}'`
- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'`
"""
# Argument sections filled into `gen:` markers on handwritten pages
engine_args = create_parser(EngineArgs.add_cli_args)
async_engine_args = create_parser(AsyncEngineArgs.add_cli_args, async_args_only=True)
fill_markers(
"configuration/engine_args.md",
{
"engine-args": (
f"{JSON_TIP}## `EngineArgs`\n\n"
f"{linkify_docs_urls(format_help(engine_args))}"
f"## `AsyncEngineArgs`\n\n"
f"{linkify_docs_urls(format_help(async_engine_args))}"
)
},
)
# CLI reference pages generated entirely from their parser: page -> (parser, JSON tip)
pages = {
"cli/serve.md": (create_parser(openai_cli_args.make_arg_parser), True),
"cli/chat.md": (create_parser(ChatCommand.add_cli_args), False),
"cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False),
"cli/run-batch.md": (create_parser(openai_run_batch.make_arg_parser), True),
"cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True),
"cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True),
# URL kept as `mm_processor` for back-compat; command name is `mm-processor`
"cli/bench/mm_processor.md": (
create_parser(BenchmarkMMProcessorSubcommand.add_cli_args),
True,
),
"cli/bench/serve.md": (create_parser(bench_serve.add_cli_args), True),
"cli/bench/startup.md": (create_parser(bench_startup.add_cli_args), True),
"cli/bench/throughput.md": (create_parser(bench_throughput.add_cli_args), True),
"cli/bench/sweep/plot.md": (create_parser(bench_sweep_plot.add_cli_args), True),
"cli/bench/sweep/plot_pareto.md": (
create_parser(bench_sweep_plot_pareto.add_cli_args),
True,
),
"cli/bench/sweep/serve.md": (create_parser(bench_sweep_serve.add_cli_args), True),
"cli/bench/sweep/serve_workload.md": (
create_parser(bench_sweep_serve_workload.add_cli_args),
True,
),
"cli/bench/sweep/startup.md": (
create_parser(bench_sweep_startup.add_cli_args),
True,
),
}
# Command name for pages whose file stem differs (URL kept for back-compat).
COMMAND_NAMES = {"cli/bench/mm_processor.md": "mm-processor"}
for doc_path, (parser, json_tip) in pages.items():
segments = Path(doc_path).relative_to("cli").with_suffix("").parts
label = COMMAND_NAMES.get(doc_path, segments[-1])
command = " ".join([*segments[:-1], label])
# `title` frontmatter keeps the nav label to just this command's segment,
# while the H1 stays the full `vllm ...` command for the page heading.
content = f"---\ntitle: {label}\n---\n\n"
content += f"# vllm {command}\n\n"
if parser.description:
content += f"## Overview\n\n{parser.description}\n\n"
# Rendered above instead of at the top of the Arguments section
parser.description = None
if json_tip:
content += JSON_TIP
content += f"## Arguments\n\n{linkify_docs_urls(format_help(parser))}"
with mkdocs_gen_files.open(doc_path, "w") as f:
f.write(content)
logger.debug("CLI reference generated: %s", doc_path)
logger.info("Total argparse docs generated: %d", len(pages) + 2)
# --- Bare subcommand (group) pages -------------------------------------------
# Mirror `vllm <group> --help`: an overview plus a table of child subcommands,
# each linked to its reference page. Children are read from the CLI registries
# so the listing can never drift from the actual subcommands. Each page is the
# `README.md` of its command directory so it becomes that section's index and is
# picked up by the existing nav globs.
import_bench_subcommands() # populate BenchmarkSubcommandBase.__subclasses__()
bench_subcommands = BenchmarkSubcommandBase.__subclasses__()
bench_children = [(cmd.name, cmd.help) for cmd in bench_subcommands]
groups = {
"cli/bench/README.md": (BenchmarkSubcommand.help, bench_children),
"cli/launch/README.md": (
launch_description,
[(cmd.name, cmd.help) for cmd in LaunchSubcommandBase.__subclasses__()],
),
"cli/bench/sweep/README.md": (
dict(bench_children).get("sweep"),
[(args.parser_name, args.parser_help) for args, _ in sweep_subcommands],
),
}
# Doc paths that exist, so we only link a child that has a reference page.
existing_pages = set(pages) | set(groups)
def child_link(group_doc: str, name: str) -> str | None:
group_dir = Path(group_doc).parent # cli/bench/README.md -> cli/bench
for stem in (name, name.replace("-", "_")):
# A leaf page (bench/latency.md) or a nested group index (sweep/README.md)
for candidate in (group_dir / f"{stem}.md", group_dir / stem / "README.md"):
if candidate.as_posix() in existing_pages:
return candidate.relative_to(group_dir).as_posix()
return None
for doc_path, (overview, children) in groups.items():
title = "vllm " + Path(doc_path).parent.relative_to("cli").as_posix()
lines = [f"# {title.replace('/', ' ')}", ""]
if overview:
lines += ["## Overview", "", overview.strip(), ""]
lines += ["## Subcommands", "", "| Command | Description |", "| --- | --- |"]
for name, summary in children:
link = child_link(doc_path, name)
command = f"[`{name}`]({link})" if link else f"`{name}`"
lines.append(f"| {command} | {(summary or '').strip()} |")
with mkdocs_gen_files.open(doc_path, "w") as f:
f.write("\n".join(lines) + "\n")
logger.debug("CLI group reference generated: %s", doc_path)
logger.info("CLI group reference pages generated: %d", len(groups))
tools/pre_commit/generate_attention_backend_docs.py → docs/mkdocs/gen_files/generate_attention_backends.py
+61 -404
View File
@@ -9,33 +9,28 @@ based on the checks in AttentionBackend.validate_configuration().
This approach avoids requiring CUDA/ROCm/GPU libraries to be installed.
When used as a pre-commit hook, this script receives filenames as arguments
and only runs the check if any of the relevant files were modified.
It runs as an mkdocs-gen-files script, so the page is generated at docs build
time rather than being committed to the repository.
"""
import argparse
import ast
import fnmatch
import logging
import sys
from collections.abc import Callable
from pathlib import Path
from typing import Any
sys.path.insert(0, str(Path(__file__).parent))
from generated_content import fill_markers # noqa: E402
logger = logging.getLogger("mkdocs")
# ---------------------------------------------------------------------------
# Constants and file paths
# ---------------------------------------------------------------------------
REPO_ROOT = Path(__file__).parent.parent.parent
RELEVANT_PATTERNS = [
"vllm/v1/attention/backends/*.py",
"vllm/v1/attention/backends/**/*.py",
"vllm/models/minimax_m3/common/sparse_attention.py",
"vllm/model_executor/layers/attention/mla_attention.py",
"vllm/platforms/cuda.py",
"tools/pre_commit/generate_attention_backend_docs.py",
"docs/design/attention_backends.md",
]
REPO_ROOT = Path(__file__).parent.parent.parent.parent
BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends"
REGISTRY_FILE = BACKENDS_DIR / "registry.py"
@@ -55,19 +50,6 @@ BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = {
}
def is_relevant_file(filepath: str) -> bool:
"""Check if a file matches any of the relevant patterns."""
path = Path(filepath)
if path.is_absolute():
try:
path = path.relative_to(REPO_ROOT)
except ValueError:
return False
path_str = str(path)
return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS)
MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill"
MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py"
MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py"
@@ -960,7 +942,7 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None
try:
tree = ast.parse(file_path.read_text())
except Exception as e:
print(f" Warning: Could not parse {file_path}: {e}", file=sys.stderr)
logger.warning("Could not parse %s: %s", file_path, e)
return None
class_name = class_path.rsplit(".", 1)[1]
@@ -1657,113 +1639,12 @@ def _render_table(
return lines
def generate_markdown_table(
backends: list[dict[str, Any]], title: str, is_mla_table: bool = False
) -> str:
"""Generate a titled markdown table from backend info."""
if not backends:
return f"## {title}\n\nNo backends found.\n"
has_versions = any(b.get("version") for b in backends)
columns = _build_columns(is_mla_table, has_versions)
lines = [f"## {title}", ""]
lines.extend(_render_table(columns, backends))
lines.append("")
return "\n".join(lines)
# ---------------------------------------------------------------------------
# Markdown section generators (usage, priority, legend, MLA)
# ---------------------------------------------------------------------------
def generate_usage_section() -> str:
"""Generate the usage documentation section."""
return """## Setting the Attention Backend
### Command Line
There are two ways to specify the backend from the command line:
**Option 1: Using `--attention-backend` (simple)**
```bash
vllm serve <model> --attention-backend FLASH_ATTN
```
**Option 2: Using `--attention-config.backend` / `-ac.backend` (structured config)**
```bash
# Dot notation
vllm serve <model> --attention-config.backend FLASH_ATTN
vllm serve <model> -ac.backend FLASH_ATTN
# JSON format
vllm serve <model> --attention-config '{"backend": "FLASH_ATTN"}'
vllm serve <model> -ac '{"backend": "FLASH_ATTN"}'
```
> **Note:** `--attention-backend` and `--attention-config.backend` are mutually
> exclusive. Use one or the other, not both.
### Python API
Use `AttentionConfig` with the `LLM` class:
```python
from vllm import LLM
from vllm.config import AttentionConfig
from vllm.v1.attention.backends.registry import AttentionBackendEnum
# Method 1: Using AttentionConfig with enum
llm = LLM(
model="Qwen/Qwen3-0.6B",
attention_config=AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN),
)
# Method 2: Using attention_backend parameter with string
llm = LLM(
model="Qwen/Qwen3-0.6B",
attention_backend="FLASH_ATTN",
)
```
## Backend Selection Behavior
### Manual Selection
When you explicitly set a backend via `--attention-backend` or `AttentionConfig`:
1. The backend is **validated** against your configuration (model dtype, head
size, compute capability, etc.)
2. If the backend **doesn't support** your configuration, an error is raised
with the specific reason
3. If valid, the backend is used
Example error when selecting an incompatible backend:
```text
ValueError: Selected backend FLASHMLA is not valid for this configuration.
Reason: ['compute capability not supported']
```
### Automatic Selection
When no backend is specified (the default):
1. vLLM iterates through backends in **priority order** (see tables below)
2. Each backend is validated against your configuration
3. The **first compatible backend** is selected
4. If no backend is compatible, an error is raised listing all backends and
their incompatibility reasons
"""
def _priority_table(
title: str,
backends: list[str],
annotations: dict[str, str] | None = None,
) -> list[str]:
"""Generate a priority table for a list of backends."""
"""Render a priority table for a list of backends."""
def _fmt(b: str) -> str:
suffix = annotations.get(b, "") if annotations else ""
@@ -1779,102 +1660,38 @@ def _priority_table(
]
def generate_priority_section(priorities: dict[str, list[str]]) -> str:
"""Generate the priority ranking section."""
lines = [
"## Backend Priority (CUDA)",
"",
"When no backend is explicitly selected, vLLM chooses the first",
"compatible backend from these priority-ordered lists.",
"",
"Priority is **1 = highest** (tried first).",
"",
"### Standard Attention (MHA, MQA, GQA)",
"",
]
sm100 = "Blackwell (SM 10.x)"
ampere = "Ampere/Hopper (SM 8.x-9.x)"
if "standard_sm100" in priorities:
lines.extend(_priority_table(sm100, priorities["standard_sm100"]))
if "standard_default" in priorities:
lines.extend(_priority_table(ampere, priorities["standard_default"]))
lines.extend(["### MLA Attention (DeepSeek-style)", ""])
mla_sm100_annotations = {
"FLASHINFER_MLA_SPARSE": "**\\***",
}
if "mla_sm100" in priorities:
lines.extend(
_priority_table(sm100, priorities["mla_sm100"], mla_sm100_annotations)
)
if "mla_default" in priorities:
lines.extend(_priority_table(ampere, priorities["mla_default"]))
if "mla_sm100" in priorities:
lines.append(
"> **\\*** For sparse MLA, FP8 KV cache always prefers "
"`FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` "
"is preferred for low query-head counts (<= 16), while "
"`FLASHMLA_SPARSE` is preferred otherwise."
)
lines.append(">")
lines.append(
"> **Note:** ROCm and CPU platforms have their own selection logic. "
"See the platform-specific documentation for details."
)
lines.append("")
return "\n".join(lines)
_SM100 = "Blackwell (SM 10.x)"
_AMPERE = "Ampere/Hopper (SM 8.x-9.x)"
def generate_legend() -> str:
"""Generate a legend explaining the table columns."""
return """## Legend
| Column | Description |
| ------ | ----------- |
| **Dtypes** | Supported model data types (fp16, bf16, fp32) |
| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) |
| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
| **Head Sizes** | Supported attention head sizes |
| **Sink** | Attention sink support (for StreamingLLM) |
| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models |
| **Sparse** | Sparse attention support (MLA only) |
| **MM Prefix** | Multimodal prefix full attention support |
| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
**Symbols:** ✅ = Supported, ❌ = Not supported
"""
def generate_mla_section(
prefill_backends: list[dict[str, Any]],
decode_backends: list[dict[str, Any]],
v4_decode_backends: list[dict[str, Any]] | None = None,
def _priority_block(
priorities: dict[str, list[str]],
sm100_key: str,
default_key: str,
sm100_annotations: dict[str, str] | None = None,
) -> str:
"""Generate the complete MLA section with prefill and decode tables."""
"""Render whichever priority tables exist for one attention category."""
lines: list[str] = []
if sm100_key in priorities:
lines += _priority_table(_SM100, priorities[sm100_key], sm100_annotations)
if default_key in priorities:
lines += _priority_table(_AMPERE, priorities[default_key])
return "\n".join(lines).strip()
def _feature_table(backends: list[dict[str, Any]], is_mla: bool) -> str:
"""Render a backend feature table (header, separator, one row per backend)."""
has_versions = any(b.get("version") for b in backends)
columns = _build_columns(is_mla, has_versions)
return "\n".join(_render_table(columns, backends))
def _mla_prefill_table(prefill_backends: list[dict[str, Any]]) -> str:
"""Render the MLA prefill backend table."""
lines = [
"## MLA (Multi-head Latent Attention) Backends",
"",
"MLA uses separate backends for prefill and decode phases.",
"",
"### Prefill Backends",
"",
"To explicitly select a prefill backend, use",
"`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).",
"Otherwise, the prefill backend is selected automatically at runtime based on",
"hardware and configuration.",
"",
"| Backend | Description | Dtypes | Compute Cap. | Notes |",
"| ------- | ----------- | ------ | ------------ | ----- |",
]
for backend in prefill_backends:
row = "| `{}`{} | {} | {} | {} | {} |".format(
backend["name"],
@@ -1885,87 +1702,21 @@ def generate_mla_section(
backend.get("notes", ""),
)
lines.append(row.replace(" ", " "))
lines.extend(
[
"",
"> **‡** Automatic selection tries FlashAttention first. On Blackwell",
"> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then",
"> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.",
"",
"### Decode Backends",
"",
"MLA decode backends are selected using the standard",
"`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).",
"",
]
)
# Reuse data-driven table rendering for decode backends
columns = _build_columns(is_mla=True, has_versions=False)
lines.extend(_render_table(columns, decode_backends))
if v4_decode_backends:
lines.extend(
[
"",
"### DeepSeek V4 Decode Backends",
"",
"DeepSeek V4 sparse MLA uses its own decode backends, selected via",
"`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,",
"`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index",
"pipeline (compressor + SWA + indexer, 256-token blocks, head 512);",
"default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and",
"`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.",
"",
]
)
lines.extend(_render_table(columns, v4_decode_backends))
lines.append("")
return "\n".join(lines)
def generate_minimax_section(backends: list[dict[str, Any]]) -> str:
"""Generate the MiniMax M3 sparse attention section."""
lines = [
"## MiniMax M3 Sparse Attention Backends",
"",
'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")',
"layers. It is wired in directly by the model and is not part of the",
"automatic priority lists above. A lightning indexer scores KV blocks, the",
"top-k blocks (plus fixed init/local blocks) are selected, and attention",
"attends only to those blocks; index keys live in a separate side cache.",
"",
]
columns = _build_columns(is_mla=False, has_versions=False)
lines.extend(_render_table(columns, backends))
lines.append("")
return "\n".join(lines)
def build_blocks() -> dict[str, str]:
"""Build the generated table blocks keyed by their `gen:` marker name.
# ---------------------------------------------------------------------------
# Top-level orchestration
# ---------------------------------------------------------------------------
def generate_docs() -> str:
"""Generate the complete documentation."""
Only the tables are generated here; the surrounding prose lives in the
handwritten ``docs/design/attention_backends.md`` page.
"""
attention_backends_map = parse_registry()
# Parse priority lists from cuda.py
priorities = parse_cuda_priority_lists()
# Parse FlashAttention FA2/FA3 feature differences
fa_features = parse_flash_attn_features()
# Parse FlashInfer TRTLLM feature differences (native vs TRTLLM on Blackwell)
fi_features = parse_flashinfer_trtllm_features()
# Parse MLA prefill backends
mla_prefill_backends = parse_mla_prefill_backends()
# Collect backend info
all_backends = []
for backend_name, class_path in attention_backends_map.items():
if backend_name in SKIP_BACKENDS:
@@ -1973,17 +1724,14 @@ def generate_docs() -> str:
info = analyze_backend(backend_name, class_path)
if info:
all_backends.append(info)
# Expand backends into version variants
if fa_features:
all_backends = _expand_flash_attn_variants(all_backends, fa_features)
if fi_features:
all_backends = _expand_flashinfer_variants(all_backends, fi_features)
# DeepSeek V4 (*_DSV4) decode backends and MiniMax M3 sparse backends each
# get their own subsection rather than mixing into the main MLA / standard
# tables (the ROCm V4 backend isn't flagged is_mla by the AST heuristic, so
# filter purely on the name).
# DeepSeek V4 (*_DSV4) and MiniMax M3 sparse backends get their own tables
# rather than mixing into the main MLA / standard tables (the ROCm V4 backend
# isn't flagged is_mla by the AST heuristic, so filter purely on the name).
def _is_v4(b: dict[str, Any]) -> bool:
return b["name"].endswith("_DSV4")
@@ -1999,112 +1747,21 @@ def generate_docs() -> str:
if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b)
]
# Generate documentation
script_path = "tools/pre_commit/generate_attention_backend_docs.py"
doc_lines = [
"# Attention Backend Feature Support",
"",
f"This document is auto-generated by `{script_path}`.",
"It shows the feature support for each registered attention backend",
"based on the checks in `AttentionBackend.validate_configuration()`.",
"",
"**Do not edit this file manually.** Run the following command to",
"regenerate it:",
"",
"```bash",
f"python {script_path}",
"```",
"",
]
# Add usage documentation
doc_lines.append(generate_usage_section())
# Add priority section
doc_lines.append(generate_priority_section(priorities))
# Add legend and feature tables
doc_lines.append(generate_legend())
standard_title = "Standard Attention (MHA, MQA, GQA) Backends"
doc_lines.append(
generate_markdown_table(non_mla_backends, standard_title, is_mla_table=False)
)
# Add footnotes for version/variant distinctions (in table order)
footnotes = []
if fi_features:
footnotes.append(
"> **†** FlashInfer Native is the regular FlashInfer path. XQA is the "
"SM90 decode path exposed through FlashInfer's TRTLLM decode API. "
"trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen "
"via `--attention-config.use_trtllm_attention=0`."
)
if fa_features:
footnotes.append(
"> **\\*** Specify the FlashAttention version via "
"`--attention-config.flash_attn_version=2`, `3`, or `4`. "
"Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), "
"FA2 otherwise."
)
if footnotes:
doc_lines.append("\n>\n".join(footnotes) + "\n")
# Add MiniMax M3 sparse section (separate category after standard GQA)
if minimax_backends:
doc_lines.append(generate_minimax_section(minimax_backends))
# Add MLA section with prefill and decode backends
doc_lines.append(
generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends)
)
return "\n".join(doc_lines)
mla_sm100_annotations = {"FLASHINFER_MLA_SPARSE": "**\\***"}
return {
"priority-standard": _priority_block(
priorities, "standard_sm100", "standard_default"
),
"priority-mla": _priority_block(
priorities, "mla_sm100", "mla_default", mla_sm100_annotations
),
"table-standard": _feature_table(non_mla_backends, is_mla=False),
"table-minimax": _feature_table(minimax_backends, is_mla=False),
"table-mla-prefill": _mla_prefill_table(mla_prefill_backends),
"table-mla-decode": _feature_table(mla_backends, is_mla=True),
"table-mla-v4-decode": _feature_table(v4_decode_backends, is_mla=True),
}
def main():
parser = argparse.ArgumentParser(
description="Generate attention backend documentation table"
)
parser.add_argument(
"--output",
"-o",
type=str,
default=str(REPO_ROOT / "docs" / "design" / "attention_backends.md"),
help="Output file path (default: docs/design/attention_backends.md)",
)
parser.add_argument(
"--check",
action="store_true",
help="Check if the documentation is up to date (for pre-commit)",
)
parser.add_argument(
"files",
nargs="*",
help="Files to check (passed by pre-commit). If none are relevant, skip.",
)
args = parser.parse_args()
if args.files and not any(is_relevant_file(f) for f in args.files):
sys.exit(0)
output_path = Path(args.output)
new_content = generate_docs()
if args.check:
needs_update = (
not output_path.exists() or output_path.read_text() != new_content
)
if needs_update:
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(new_content)
print(f"🔄 Regenerated: {output_path}")
sys.exit(1)
print(f"✅ Up to date: {output_path}")
sys.exit(0)
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(new_content)
print(f"Generated: {output_path}")
if __name__ == "__main__":
main()
logger.info("Generating attention backend documentation")
fill_markers("design/attention_backends.md", build_blocks())
docs/mkdocs/hooks/generate_examples.py → docs/mkdocs/gen_files/generate_examples.py
+34 -41
View File
@@ -5,16 +5,15 @@ import logging
from dataclasses import dataclass
from functools import cached_property
from pathlib import Path
from typing import Literal
import mkdocs_awesome_nav.nav.directory as _nav_dir
import mkdocs_gen_files
import regex as re
logger = logging.getLogger("mkdocs")
ROOT_DIR = Path(__file__).parent.parent.parent.parent
ROOT_DIR_RELATIVE = "../../../../.."
EXAMPLE_DIR = ROOT_DIR / "examples"
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/examples"
def title(text: str) -> str:
@@ -197,44 +196,38 @@ class Example:
return content
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
# Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
# title-cased (e.g. "Offline Inference" instead of "Offline inference").
import mkdocs_awesome_nav.nav.directory as _nav_dir
# Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
# title-cased (e.g. "Offline Inference" instead of "Offline inference").
_nav_dir.dirname_to_title = title
logger.info("Generating example documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
_nav_dir.dirname_to_title = title
logger.info("Generating example documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
categories = sorted(
p for p in EXAMPLE_DIR.iterdir() if p.is_dir() and not p.name.startswith(".")
)
# Create the EXAMPLE_DOC_DIR if it doesn't exist
if not EXAMPLE_DOC_DIR.exists():
EXAMPLE_DOC_DIR.mkdir(parents=True)
examples = []
glob_patterns = ["*.py", "*.md", "*.sh"]
# Find categorised examples
for category in categories:
logger.info("Processing category: %s", category.stem)
globs = [category.glob(pattern) for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path, category.stem))
# Find examples in subdirectories
globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path.parent, category.stem))
categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
examples = []
glob_patterns = ["*.py", "*.md", "*.sh"]
# Find categorised examples
for category in categories:
logger.info("Processing category: %s", category.stem)
globs = [category.glob(pattern) for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path, category.stem))
# Find examples in subdirectories
globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path.parent, category.stem))
# Generate the example documentation
for example in sorted(examples, key=lambda e: e.path.stem):
example_name = f"{example.path.stem}.md"
doc_path = EXAMPLE_DOC_DIR / example.category / example_name
if not doc_path.parent.exists():
doc_path.parent.mkdir(parents=True)
# Specify encoding for building on Windows
with open(doc_path, "w+", encoding="utf-8") as f:
f.write(example.generate())
logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))
logger.info("Total examples generated: %d", len(examples))
# Generate the example documentation
for example in sorted(examples, key=lambda e: e.path.stem):
doc_path = f"examples/{example.category}/{example.path.stem}.md"
with mkdocs_gen_files.open(doc_path, "w") as f:
f.write(example.generate())
if example.main_file is not None:
# Point the edit button at the example's source file
edit_path = Path("..") / example.main_file.relative_to(ROOT_DIR)
mkdocs_gen_files.set_edit_path(doc_path, str(edit_path))
logger.debug("Example generated: %s", doc_path)
logger.info("Total examples generated: %d", len(examples))
docs/mkdocs/hooks/generate_metrics.py → docs/mkdocs/gen_files/generate_metrics.py
+31 -44
View File
@@ -2,27 +2,28 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import ast
import logging
import sys
from pathlib import Path
from typing import Literal
sys.path.insert(0, str(Path(__file__).parent))
from generated_content import fill_markers # noqa: E402
logger = logging.getLogger("mkdocs")
ROOT_DIR = Path(__file__).parent.parent.parent.parent
DOCS_DIR = ROOT_DIR / "docs"
GENERATED_METRICS_DIR = DOCS_DIR / "generated" / "metrics"
# Files to scan for metric definitions - each will generate a separate table
# Files to scan for metric definitions - each fills a `gen:` marker in
# docs/usage/metrics.md with its table (the section heading and any preamble
# live in the tracked page next to the marker).
METRIC_SOURCE_FILES = [
{"path": "vllm/v1/metrics/loggers.py", "output": "general.inc.md"},
{
"path": "vllm/v1/spec_decode/metrics.py",
"output": "spec_decode.inc.md",
},
{"path": "vllm/v1/metrics/loggers.py", "key": "metrics-general"},
{"path": "vllm/v1/spec_decode/metrics.py", "key": "metrics-spec-decode"},
{
"path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py",
"output": "nixl_connector.inc.md",
"key": "metrics-nixl",
},
{"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"},
{"path": "vllm/v1/metrics/perf.py", "key": "metrics-mfu"},
]
@@ -110,41 +111,27 @@ def generate_markdown_table(metrics: list[dict[str, str]]) -> str:
return "\n".join(lines) + "\n"
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
"""Generate metrics documentation tables from source files."""
logger.info("Generating metrics documentation")
logger.info("Generating metrics documentation")
# Create generated directory if it doesn't exist
GENERATED_METRICS_DIR.mkdir(parents=True, exist_ok=True)
blocks = {}
total_metrics = 0
for source_config in METRIC_SOURCE_FILES:
source_path = source_config["path"]
total_metrics = 0
for source_config in METRIC_SOURCE_FILES:
source_path = source_config["path"]
output_file = source_config["output"]
filepath = ROOT_DIR / source_path
if not filepath.exists():
raise FileNotFoundError(f"Metrics source file not found: {filepath}")
filepath = ROOT_DIR / source_path
if not filepath.exists():
raise FileNotFoundError(f"Metrics source file not found: {filepath}")
logger.debug("Extracting metrics from: %s", source_path)
metrics = extract_metrics_from_file(filepath)
logger.debug("Found %d metrics in %s", len(metrics), source_path)
logger.debug("Extracting metrics from: %s", source_path)
metrics = extract_metrics_from_file(filepath)
logger.debug("Found %d metrics in %s", len(metrics), source_path)
blocks[source_config["key"]] = generate_markdown_table(metrics).strip()
total_metrics += len(metrics)
# Generate and write the markdown table for this source
table_content = generate_markdown_table(metrics)
output_path = GENERATED_METRICS_DIR / output_file
with open(output_path, "w", encoding="utf-8") as f:
f.write(table_content)
total_metrics += len(metrics)
logger.info(
"Generated metrics table: %s (%d metrics)",
output_path.relative_to(ROOT_DIR),
len(metrics),
)
logger.info(
"Total metrics generated: %d across %d files",
total_metrics,
len(METRIC_SOURCE_FILES),
)
fill_markers("usage/metrics.md", blocks)
logger.info(
"Total metrics generated: %d across %d files",
total_metrics,
len(METRIC_SOURCE_FILES),
)
docs/mkdocs/gen_files/generated_content.py
+56
View File
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Inline build-time generated content into existing docs pages.
Source pages mark where generated content goes with a snippet-style marker,
`--8<-- "gen:<key>"`, so the insertion point is explicit and readable. The
substitution happens here (at gen-files time, before mkdocs-gen-files shadows
the page), not via pymdownx.snippets, so the content can be generated at build
time without living in a real file on disk.
The `gen:` prefix keeps these markers distinct from real pymdownx.snippets
includes, and `fill_markers` fails loudly if a marker is missing or left behind
(pymdownx.snippets would otherwise silently drop an unsubstituted marker).
"""
from pathlib import Path
import mkdocs_gen_files
import regex as re
DOCS_DIR = Path(__file__).parent.parent.parent
_MARKER = '--8<-- "gen:{key}"'
_ANY_MARKER = re.compile(r'--8<-- "gen:[^"]*"')
def fill_markers(doc_path: str, blocks: dict[str, str]) -> None:
"""Replace `--8<-- "gen:<key>"` markers in a docs page with generated content.
Args:
doc_path: Docs-relative path of the source page to fill.
blocks: Mapping of marker key to the markdown to insert in its place.
Raises:
FileNotFoundError: If the source page does not exist.
ValueError: If an expected marker is missing, or any `gen:` marker is
left unsubstituted after filling.
"""
source = DOCS_DIR / doc_path
if not source.exists():
raise FileNotFoundError(f"Cannot fill markers in missing page: {doc_path}")
text = source.read_text()
for key, content in blocks.items():
marker = _MARKER.format(key=key)
if marker not in text:
raise ValueError(f"{doc_path}: missing marker {marker}")
text = text.replace(marker, content)
if leftover := _ANY_MARKER.search(text):
raise ValueError(f"{doc_path}: unsubstituted marker {leftover.group()}")
with mkdocs_gen_files.open(doc_path, "w") as f:
f.write(text)
# Keep the edit button pointing at the real source page
mkdocs_gen_files.set_edit_path(doc_path, doc_path)
docs/mkdocs/hooks/url_schemes.py
+39 -4
View File
@@ -19,6 +19,7 @@ The on_page_markdown hook passes the current page context to the preprocessor be
each page is converted.
"""
import posixpath
from pathlib import Path
import regex as re
@@ -38,18 +39,22 @@ TITLE = r"(?P<title>[^\[\]<>]+?)"
REPO = r"(?P<repo>.+?/.+?)"
TYPE = r"(?P<type>issues|pull|projects)"
NUMBER = r"(?P<number>\d+)"
VERSION = r"[^/\s]+"
PATH = r"(?P<path>[^\s]+?)"
FRAGMENT = r"(?P<fragment>#[^\s]+)?"
URL = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}"
URL_GITHUB = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}"
RELATIVE = rf"(?!(https?|ftp)://|#){PATH}{FRAGMENT}"
URL_DOCS = f"https://docs.vllm.ai/en/{VERSION}/{PATH}{FRAGMENT}"
# Common titles to use for GitHub links when none is provided in the link.
TITLES = {"issues": "Issue ", "pull": "Pull Request ", "projects": "Project "}
# Regex to match GitHub issue, PR, and project links with optional titles.
github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL}(\)|>)")
github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL_GITHUB}(\)|>)")
# Regex to match relative file links with optional titles.
relative_link = re.compile(rf"\[{TITLE}\]\({RELATIVE}\)")
# Regex to match absolute docs.vllm.ai links (should only exist in CLI).
docs_link = re.compile(rf"\[{TITLE}\]\({URL_DOCS}\)")
class UrlSchemesPreprocessor(Preprocessor):
@@ -61,7 +66,8 @@ class UrlSchemesPreprocessor(Preprocessor):
def run(self, lines):
page = self.ext.page
if page is None or getattr(page.file, "abs_src_path", None) is None:
files = self.ext.files
if page is None:
return lines
def replace_relative_link(match: re.Match) -> str:
@@ -70,7 +76,7 @@ class UrlSchemesPreprocessor(Preprocessor):
"""
title = match.group("title")
path = match.group("path")
path = (Path(page.file.abs_src_path).parent / path).resolve()
path = ((DOC_DIR / page.file.src_uri).parent / path).resolve()
fragment = match.group("fragment") or ""
# Check if the path exists and is outside the docs dir
@@ -105,9 +111,36 @@ class UrlSchemesPreprocessor(Preprocessor):
url = f"https://github.com/{repo}/{type}/{number}{fragment}"
return f"[{gh_icon} {title}]({url})"
def replace_docs_link(match: re.Match) -> str:
"""Rewrite absolute docs.vllm.ai links as doc-relative links."""
title = match.group("title")
path = match.group("path").rstrip("/")
fragment = match.group("fragment") or ""
# vllm.config.<Class> API reference -> mkdocstrings cross-reference
if path == "api/vllm/config" and re.fullmatch(
r"#vllm\.config\.\w+", fragment
):
ident = fragment[1:]
return f"[`{ident}`][{ident}]"
# Other docs pages -> link relative to the current page, but only
# when the target is a known docs page (real or generated); leave
# unknown/external URLs untouched. This is correct even when the same
# docstring is also rendered on its API reference page.
src = f"{path.removesuffix('.html')}.md"
if files.get_file_from_path(src) is None:
return match.group(0)
rel = posixpath.relpath(src, posixpath.dirname(page.file.src_uri))
# Auto-wrapped bare URLs use the URL as their title; make it readable.
if title.startswith("http"):
title = path.removesuffix(".html")
return f"[{title}]({rel}{fragment})"
markdown = "\n".join(lines)
markdown = github_link.sub(replace_github_link, markdown)
markdown = relative_link.sub(replace_relative_link, markdown)
markdown = docs_link.sub(replace_docs_link, markdown)
return markdown.split("\n")
@@ -116,6 +149,7 @@ class UrlSchemesExtension(Extension):
def __init__(self, **kwargs):
self.page = None
self.files = None
super().__init__(**kwargs)
def extendMarkdown(self, md):
@@ -138,4 +172,5 @@ def on_page_markdown(
) -> str:
"""Pass the current page context to the preprocessor."""
_ext.page = page
_ext.files = files
return markdown
docs/models/supported_models.md
+4 -4
View File
@@ -19,7 +19,7 @@ vLLM also supports model implementations that are available in Transformers. We
Currently, the Transformers modeling backend works for the following:
- Modalities: embedding models, language models and vision-language models*
- Modalities: embedding models, language models, vision-language models* and audio-language models
- Architectures: encoder-only, decoder-only, mixture-of-experts
- Attention types: full attention and/or sliding attention
@@ -427,7 +427,6 @@ th {
| `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |
| `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ |
| `OuroForCausalLM` | ouro | `ByteDance/Ouro-1.4B`, `ByteDance/Ouro-2.6B`, etc. | ✅︎ | |
| `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |
| `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
| `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
@@ -435,7 +434,6 @@ th {
| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
| `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ |
| `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ |
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ |
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ |
@@ -466,6 +464,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
| `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ |
!!! note
Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
@@ -608,7 +607,8 @@ Some models are supported only via the [Transformers modeling backend](#transfor
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
| ------------ | ------ | ------ | ----------------- | --------------------------- | --------------------------------------- |
| `Emu3ForConditionalGeneration` | Emu3 | T + I | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ |
| `Emu3ForConditionalGeneration` | Emu3 | T + I<sup>+</sup> | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ |
| `VibeVoiceAsrForConditionalGeneration` | VibeVoice-ASR | T + A<sup>+</sup> | `microsoft/VibeVoice-ASR-HF` | ✅︎ | ✅︎ |
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
docs/pre_run_check.sh
+1 -1
View File
@@ -28,7 +28,7 @@ DOCS_PATHS=(
docs/ # Actual docs content
examples/ # Examples are rendered in docs
vllm/ # API & CLI reference
requirements/test/cuda.txt # CLI reference (see docs/mkdocs/hooks/generate_argparse.py)
requirements/test/cuda.txt # CLI reference (see docs/mkdocs/gen_files/generate_argparse.py)
mkdocs.yaml # Affects build process
.readthedocs.yaml # Affects build process
requirements/docs.txt # Affects build process

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